fix: add support for autodiff through VoA broadcast

Author: AayushSabharwal

ext/RecursiveArrayToolsZygoteExt.jl

@@ -11,7 +11,7 @@ else
 end
 
 # Define a new species of projection operator for this type:
-ChainRulesCore.ProjectTo(x::VectorOfArray) = ChainRulesCore.ProjectTo{VectorOfArray}()
+# ChainRulesCore.ProjectTo(x::VectorOfArray) = ChainRulesCore.ProjectTo{VectorOfArray}()
 
 function ChainRulesCore.rrule(T::Type{<:RecursiveArrayTools.GPUArraysCore.AbstractGPUArray},
     xs::AbstractVectorOfArray)
@@ -117,4 +117,53 @@ end
     A.x, literal_ArrayPartition_x_adjoint
 end
 
+@adjoint function Array(VA::AbstractVectorOfArray)
+    Array(VA),
+    y -> (Array(y),)
 end
+
+
+ChainRulesCore.ProjectTo(a::AbstractVectorOfArray) = ChainRulesCore.ProjectTo{VectorOfArray}((sz = size(a)))
+
+function (p::ChainRulesCore.ProjectTo{VectorOfArray})(x)
+    arr = reshape(x, p.sz)
+    return VectorOfArray([arr[:, i] for i in 1:p.sz[end]])
+end
+
+function Zygote.unbroadcast(x::AbstractVectorOfArray, x̄)
+    N = ndims(x̄)
+    if length(x) == length(x̄)
+        Zygote._project(x, x̄)  # ProjectTo handles reshape, offsets, structured matrices, row vectors
+    else
+        dims = ntuple(d -> size(x, d) == 1 ? d : ndims(x̄)+1, ndims(x̄))
+        Zygote._project(x, Zygote.accum_sum(x̄; dims = dims))
+    end
+end
+
+@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::F, a::AbstractVectorOfArray, b) where {F} = _broadcast_generic(__context__, f, a, b)
+@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::F, a, b::AbstractVectorOfArray) where {F} = _broadcast_generic(__context__, f, a, b)
+@adjoint Broadcast.broadcasted(::Broadcast.AbstractArrayStyle, f::F, a::AbstractVectorOfArray, b::AbstractVectorOfArray) where {F} = _broadcast_generic(__context__, f, a, b)
+
+@inline function _broadcast_generic(__context__, f::F, args...) where {F}
+    T = Broadcast.combine_eltypes(f, args)
+    # Avoid generic broadcasting in two easy cases:
+    if T == Bool
+        return (f.(args...), _ -> nothing)
+    elseif T <: Union{Real, Complex} && isconcretetype(T) && Zygote._dual_purefun(F) && all(Zygote._dual_safearg, args) && !Zygote.isderiving()
+      return Zygote.broadcast_forward(f, args...)
+    end
+    len = Zygote.inclen(args)
+    y∂b = Zygote._broadcast((x...) -> Zygote._pullback(__context__, f, x...), args...)
+    y = broadcast(first, y∂b)
+    function ∇broadcasted(ȳ)
+        y∂b = y∂b isa AbstractVectorOfArray ? Iterators.flatten(y∂b.u) : y∂b
+        ȳ = ȳ isa AbstractVectorOfArray ? Iterators.flatten(ȳ.u) : ȳ
+        dxs_zip = map(((_, pb), ȳ₁) -> pb(ȳ₁), y∂b, ȳ)
+        getters = ntuple(i -> Zygote.StaticGetter{i}(), len)
+        dxs = map(g -> Zygote.collapse_nothings(map(g, dxs_zip)), getters)
+        (nothing, Zygote.accum_sum(dxs[1]), map(Zygote.unbroadcast, args, Base.tail(dxs))...)
+    end
+    return y, ∇broadcasted
+end
+
+end # module