fix up some of the VectorOfArray dispatches

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "RecursiveArrayTools"
 uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "2.18.0"
+version = "2.19.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/zygote.jl

@@ -54,22 +54,48 @@ ChainRulesCore.ProjectTo(x::VectorOfArray) = ChainRulesCore.ProjectTo{VectorOfAr
 # These rules duplicate the `rrule` methods above, because Zygote looks for an `@adjoint`
 # definition first, and finds its own before finding those.
 
-ZygoteRules.@adjoint function getindex(VA::AbstractVectorOfArray, i::Union{Int,AbstractArray{Int},CartesianIndex,Colon,BitArray,AbstractArray{Bool}})
+ZygoteRules.@adjoint function getindex(VA::AbstractVectorOfArray, i::Int)
   function AbstractVectorOfArray_getindex_adjoint(Δ)
     Δ′ = [(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}}...)
+ZygoteRules.@adjoint function getindex(VA::AbstractVectorOfArray, i::Union{BitArray,AbstractArray{Bool}})
+  function AbstractVectorOfArray_getindex_adjoint(Δ)
+    Δ′ = [(i[j] ? Δ[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::AbstractArray{Int})
+  function AbstractVectorOfArray_getindex_adjoint(Δ)
+    iter = 0
+    Δ′ = [(j ∈ i ? Δ[iter+=1] : 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}})
+  function AbstractVectorOfArray_getindex_adjoint(Δ)
+    Δ′ = [(i[j] ? Δ[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::Int, j::Union{Int,AbstractArray{Int},CartesianIndex,Colon,BitArray,AbstractArray{Bool}}...)
   function AbstractVectorOfArray_getindex_adjoint(Δ)
     Δ′ = [(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
 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)