fix: view adjoints

SciML · Jan 22, 2024 · 2af86b7 · 2af86b7
1 parent cc55ee6
commit 2af86b7
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/ext/RecursiveArrayToolsZygoteExt.jl b/ext/RecursiveArrayToolsZygoteExt.jl
@@ -118,11 +118,19 @@ end
     Array(VA), adj
 end
 
+@adjoint function Base.view(A::AbstractVectorOfArray, I::Colon...)
+    function adjoint(y)
+        (recursivecopy(A), map(_ -> nothing, I)...)
+    end
+    return view(A, I...), adjoint
+end
+
 @adjoint function Base.view(A::AbstractVectorOfArray, I...)
     adj = let A = A, I = I
         function view_adjoint(y)
-            A = zero(A)
-            view(A, I...) .= y
+            A = recursivecopy(A)
+            recursivefill!(A, zero(eltype(A)))
+            A[I...] .= y
             return (A, map(_ -> nothing, I)...)
         end
     end