Weird allocation issue for fill! with ArrayPartition in Julia < 1.6

Ok, I figured out a workaround:

Even though one can iterate over a eachindex of an ArrayPartition, it is probably better to “unpack” its elements when doing this, since an ArrayPartition can contain arrays of different types.

This seems to work:

using RecursiveArrayTools,BenchmarkTools

function setZero!(PartArr::ArrayPartition)
    for arr in PartArr.x
        fill!(arr,0.)
    end
end
    
function test2()
    PartArray = ArrayPartition(
        zeros(2),
        zeros(200,200,30), 
        zeros(40,40,20,30),
    )
    @btime fill!($PartArray,0.)  #old version
    @btime setZero!($PartArray)
end

julia> test2()
  72.663 ms (4318980 allocations: 65.90 MiB) #using fill!
  1.574 ms (0 allocations: 0 bytes) #using setzero!

Its should also be possible to overload Base.fill! for an ArrayPartition to allow the efficient use of fill!.