fix: remove linear indexing from getindex adjoints

Author: AayushSabharwal

ext/RecursiveArrayToolsZygoteExt.jl

@@ -24,7 +24,7 @@ end
               for (x, j) in zip(VA.u, 1:length(VA))]
         (VectorOfArray(Δ′), nothing)
     end
-    VA[i], AbstractVectorOfArray_getindex_adjoint
+    VA[:, i], AbstractVectorOfArray_getindex_adjoint
 end
 
 @adjoint function getindex(VA::AbstractVectorOfArray,
@@ -34,7 +34,7 @@ end
               for (x, j) in zip(VA.u, 1:length(VA))]
         (VectorOfArray(Δ′), nothing)
     end
-    VA[i], AbstractVectorOfArray_getindex_adjoint
+    VA[:, i], AbstractVectorOfArray_getindex_adjoint
 end
 
 @adjoint function getindex(VA::AbstractVectorOfArray, i::AbstractArray{Int})
@@ -44,24 +44,24 @@ end
               for (x, j) in zip(VA.u, 1:length(VA))]
         (VectorOfArray(Δ′), nothing)
     end
-    VA[i], AbstractVectorOfArray_getindex_adjoint
+    VA[:, i], AbstractVectorOfArray_getindex_adjoint
 end
 
 @adjoint function getindex(VA::AbstractVectorOfArray,
-    i::Union{Int, AbstractArray{Int}})
+    i::Int)
     function AbstractVectorOfArray_getindex_adjoint(Δ)
-        Δ′ = [(i[j] ? Δ[j] : FillArrays.Fill(zero(eltype(x)), size(x)))
+        Δ′ = [(i == j ? Δ[j] : FillArrays.Fill(zero(eltype(x)), size(x)))
               for (x, j) in zip(VA.u, 1:length(VA))]
         (VectorOfArray(Δ′), nothing)
     end
-    VA[i], AbstractVectorOfArray_getindex_adjoint
+    VA[:, i], AbstractVectorOfArray_getindex_adjoint
 end
 
 @adjoint function getindex(VA::AbstractVectorOfArray, i::Colon)
     function AbstractVectorOfArray_getindex_adjoint(Δ)
         (VectorOfArray(Δ), nothing)
     end
-    VA[i], AbstractVectorOfArray_getindex_adjoint
+    VA.u[i], AbstractVectorOfArray_getindex_adjoint
 end
 
 @adjoint function getindex(VA::AbstractVectorOfArray, i::Int,