hotfix add back some zygoterules

Author: ChrisRackauckas

Project.toml

@@ -12,6 +12,7 @@ RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 ArrayInterface = "2.7, 3.0"

src/RecursiveArrayTools.jl

@@ -10,6 +10,7 @@ using DocStringExtensions
 
   import ChainRulesCore
   import ChainRulesCore: NoTangent
+  import ZygoteRules
   abstract type AbstractVectorOfArray{T, N, A} <: AbstractArray{T, N} end
   abstract type AbstractDiffEqArray{T, N, A} <: AbstractVectorOfArray{T, N, A} end
 

src/zygote.jl

@@ -42,3 +42,45 @@ function ChainRulesCore.rrule(::typeof(getproperty),A::ArrayPartition, s::Symbol
     end
     A.x,literal_ArrayPartition_x_adjoint
 end
+
+ZygoteRules.@adjoint function getindex(VA::AbstractVectorOfArray, i)
+  function AbstractVectorOfArray_getindex_adjoint(Δ)
+    Δ′ = [ (i == j ? Δ : zero(x)) for (x,j) in zip(VA.u, 1:length(VA))]
+    (Δ′,nothing)
+  end
+  VA[i],AbstractVectorOfArray_getindex_adjoint
+end
+
+ZygoteRules.@adjoint function getindex(VA::AbstractVectorOfArray, i, j...)
+  function AbstractVectorOfArray_getindex_adjoint(Δ)
+    Δ′ = zero(VA)
+    Δ′[i,j...] = Δ
+    (Δ′, i,map(_ -> nothing, j)...)
+  end
+  VA[i,j...],AbstractVectorOfArray_getindex_adjoint
+end
+
+ZygoteRules.@adjoint function ArrayPartition(x::S, ::Type{Val{copy_x}} = Val{false}) where {S<:Tuple,copy_x}
+  function ArrayPartition_adjoint(_y)
+      y = Array(_y)
+      starts = vcat(0,cumsum(reduce(vcat,length.(x))))
+      ntuple(i -> reshape(y[starts[i]+1:starts[i+1]], size(x[i])), length(x)), nothing
+  end
+
+  ArrayPartition(x, Val{copy_x}), ArrayPartition_adjoint
+end
+
+ZygoteRules.@adjoint function VectorOfArray(u)
+  VectorOfArray(u),y -> ([y[ntuple(x->Colon(),ndims(y)-1)...,i] for i in 1:size(y)[end]],)
+end
+
+ZygoteRules.@adjoint function DiffEqArray(u,t)
+  DiffEqArray(u,t),y -> ([y[ntuple(x->Colon(),ndims(y)-1)...,i] for i in 1:size(y)[end]],nothing)
+end
+
+ZygoteRules.@adjoint function ZygoteRules.literal_getproperty(A::ArrayPartition, ::Val{:x})
+  function literal_ArrayPartition_x_adjoint(d)
+      (ArrayPartition((isnothing(d[i]) ? zero(A.x[i]) : d[i] for i in 1:length(d))...),)
+  end
+  A.x,literal_ArrayPartition_x_adjoint
+end