fix pullbacks and use some FillArrays

Author: ChrisRackauckas

Project.toml

@@ -7,6 +7,7 @@ version = "2.17.2"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"

src/RecursiveArrayTools.jl

@@ -11,6 +11,9 @@ using Requires, RecipesBase, StaticArrays, Statistics,
 import ChainRulesCore
 import ChainRulesCore: NoTangent
 import ZygoteRules
+
+using FillArrays
+
 abstract type AbstractVectorOfArray{T, N, A} <: AbstractArray{T, N} end
 abstract type AbstractDiffEqArray{T, N, A} <: AbstractVectorOfArray{T, N, A} end
 

src/zygote.jl

@@ -47,7 +47,7 @@ end
 ChainRulesCore.ProjectTo(x::VectorOfArray) = ChainRulesCore.ProjectTo{VectorOfArray}()
 
 # Gradient from iteration will be e.g. Vector{Vector}, this makes it another AbstractMatrix
-(::ChainRulesCore.ProjectTo{VectorOfArray})(dx::AbstractVector{<:AbstractArray}) = VectorOfArray(dx)
+#(::ChainRulesCore.ProjectTo{VectorOfArray})(dx::AbstractVector{<:AbstractArray}) = VectorOfArray(dx)
 # Gradient from broadcasting will be another AbstractArray
 #(::ChainRulesCore.ProjectTo{VectorOfArray})(dx::AbstractArray) = dx
 
@@ -56,16 +56,16 @@ ChainRulesCore.ProjectTo(x::VectorOfArray) = ChainRulesCore.ProjectTo{VectorOfAr
 
 ZygoteRules.@adjoint function getindex(VA::AbstractVectorOfArray, i::Union{Int,AbstractArray{Int},CartesianIndex,Colon,BitArray,AbstractArray{Bool}})
   function AbstractVectorOfArray_getindex_adjoint(Δ)
-    Δ′ = [ (i == j ? Δ : zero(x)) for (x,j) in zip(VA.u, 1:length(VA))]
+    Δ′ = [(i == j ? Δ : Fill(zero(eltype(x)),size(x))) for (x,j) in zip(VA.u, 1:length(VA))]
     (VectorOfArray(Δ′),nothing)
   end
   VA[i],AbstractVectorOfArray_getindex_adjoint
 end
 
 ZygoteRules.@adjoint function getindex(VA::AbstractVectorOfArray, i::Union{Int,AbstractArray{Int},CartesianIndex,Colon,BitArray,AbstractArray{Bool}}, j::Union{Int,AbstractArray{Int},CartesianIndex,Colon,BitArray,AbstractArray{Bool}}...)
   function AbstractVectorOfArray_getindex_adjoint(Δ)
-    Δ′ = zero(VA)
-    Δ′[i,j...] = Δ
+    Δ′ = [(i == j ? zero(x) : Fill(zero(eltype(x)),size(x))) for (x,j) in zip(VA.u, 1:length(VA))]
+    Δ′[i][j...] = Δ
     (VectorOfArray(Δ′), nothing, map(_ -> nothing, j)...)
   end
   VA[i,j...],AbstractVectorOfArray_getindex_adjoint