feat: allow reverse pass over DiffEqArray as well

Author: DhairyaLGandhi

ext/RecursiveArrayToolsZygoteExt.jl

@@ -100,16 +100,29 @@ end
 end
 
 Zygote.@adjoint function Zygote.literal_getproperty(A::RecursiveArrayTools.AbstractVectorOfArray, ::Val{:u})
-    function literal_VectorOfArray_x_adjoint(d)
-        m = map(enumerate(d)) do (idx, d_i)
-            isnothing(d_i) && return zero(A.u[idx])
-            d_i
-        end
-        (VectorOfArray(m), nothing)
+    function literal_AbstractVofA_u_adjoint(d)
+        dA = vofa_u_adjoint(d, A)
+        (dA, nothing)
     end
     A.u, literal_VectorOfArray_x_adjoint
 end
 
+function vofa_u_adjoint(d, A::RecursiveArrayTools.VectorOfArray)
+    m = map(enumerate(d)) do (idx, d_i)
+        isnothing(d_i) && return zero(A.u[idx])
+        d_i 
+    end 
+    VectorOfArray(m)
+end
+
+function vofa_u_adjoint(d, A::RecursiveArrayTools.DiffEqArray)
+    m = map(enumerate(d)) do (idx, d_i)
+        isnothing(d_i) && return zero(A.u[idx])
+        d_i
+    end
+    DiffEqArray(m, A.t)
+end
+
 @adjoint function 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))...),)