fixup! fix: fix view adjoints

Author: AayushSabharwal

Project.toml

@@ -4,6 +4,7 @@ authors = ["Chris Rackauckas <[email protected]>"]
 version = "3.19.0"
 
 [deps]
+Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -36,6 +37,7 @@ RecursiveArrayToolsTrackerExt = "Tracker"
 RecursiveArrayToolsZygoteExt = "Zygote"
 
 [compat]
+Accessors = "0.1"
 Adapt = "3.4, 4"
 Aqua = "0.8"
 ArrayInterface = "7.6"

ext/RecursiveArrayToolsZygoteExt.jl

@@ -1,6 +1,7 @@
 module RecursiveArrayToolsZygoteExt
 
 using RecursiveArrayTools
+using RecursiveArrayTools.Accessors: @set, @reset
 
 if isdefined(Base, :get_extension)
     using Zygote
@@ -125,9 +126,10 @@ end
 @adjoint function Base.view(A::AbstractVectorOfArray, I::Colon...)
     view_adjoint = let A = A, I = I
         function (y)
-            A = recursivecopy(A)
-            A .= y
-            (A, map(_ -> nothing, I)...)
+            u = collect.(eachslice(y, dims=ndims(y)))
+            B = @set A.u = u
+            
+            (B, map(_ -> nothing, I)...)
         end
     end
     return view(A, I...), view_adjoint
@@ -136,11 +138,9 @@ end
 @adjoint function Base.view(A::AbstractVectorOfArray, I...)
     view_adjoint = let A = A, I = I
         function (y)
-            A = recursivecopy(A)
-            recursivefill!(A, zero(eltype(A)))
-            v = view(A, I...)
-            v .= y
-            return (A, map(_ -> nothing, I)...)
+            B = @set A .= zero(eltype(A))
+            @reset B[I...] = y
+            return (B, map(_ -> nothing, I)...)
         end
     end
     view(A, I...), view_adjoint

src/RecursiveArrayTools.jl

@@ -8,6 +8,7 @@ using DocStringExtensions
 using RecipesBase, StaticArraysCore, Statistics,
       ArrayInterface, LinearAlgebra
 using SymbolicIndexingInterface
+import Accessors
 using SparseArrays
 
 import Adapt