Type instability when iterating nested recursive arrays

[lukas-weber]

I noticed that the following is type unstable (on v3.3.3 and julia 1.10.0)

using RecursiveArrayTools

sum(ArrayPartition(ArrayPartition(zeros(4,4)))) # ::Any

sum(VectorOfArray([zeros(4), VectorOfArray([zeros(4,4)])])) # ::Any

Curiously, all of the following are stable.

sum(ArrayPartition(ArrayPartition(zeros(4)))) # ::Float64
sum(ArrayPartition(zeros(4,4))) # ::Float64

sum(VectorOfArray([VectorOfArray([zeros(4,4)])])) # ::Float64

The other recursive array types seem to be affected as well. Not sure if getting these stable is in-scope, but it would be super helpful.

[ChrisRackauckas]

Yeah getting this stable is in scope. Thanks for bringing this up. I think this may have regressed with the v3.

@AayushSabharwal

[AayushSabharwal]

VectorOfArray([zeros(4), VectorOfArray([zeros(4,4)])])

this is a bit of a tough one. Technically it isn't a valid VoA since VectorOfArray isn't an array. The eltype of this is Any, so it infers `Any.

[AayushSabharwal]

I've fixed the eltype, but sum/mapreduce isn't going to be type-stable. sum(VectorOfArray([zeros(4), VectorOfArray([zeros(4,4)])])) didn't even work in v2. I've made it work (in some cases) here, but don't think it'll ever be inferrable. As long as VectorOfArrays don't contain both arrays and VectorOfArrays, it will infer fine

[lukas-weber]

I see, makes sense. Thanks a lot!

[AayushSabharwal]

So there's a tradeoff. If mapreduce is defined as reduce(op, (mapreduce(f, op, x) for x in A.x)) then sum(ArrayPartition(ArrayPartition(zeros(4,4)))) is not type-stable (infers Any) but sum(ArrayPartition([1, 2], [3.0, 4.0])) is (infers Float64, and internally unrolls the loop). However, if it is defined as mapreduce(f, op, (x for x in A)) then the former is correctly inferred as Float64 but the latter infers Union{Int, Float64} and doesn't unroll loops.

[ChrisRackauckas]

Just check if the eltype of the tuple is ArrayPartition? That won't hit every case but it should hit "most"

[ChrisRackauckas]

This was closed but was it solved? Is https://github.com/SciML/SciMLSensitivity.jl/actions/runs/7433580549/job/20226765424?pr=965 not related?