fix: GPU tests, CuArray conversion, autodiff

SciML · Dec 28, 2023 · c1249b3 · c1249b3
1 parent 87ef7d5
commit c1249b3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/ext/RecursiveArrayToolsZygoteExt.jl b/ext/RecursiveArrayToolsZygoteExt.jl
@@ -95,17 +95,17 @@ end
     VectorOfArray(u),
     y -> begin
         y isa Ref && (y = VectorOfArray(y[].u))
-        (VectorOfArray([y[ntuple(x -> Colon(), ndims(y.u) - 1)..., i]
-                        for i in 1:size(y.u)[end]]),)
+        (VectorOfArray([y[ntuple(x -> Colon(), ndims(y) - 1)..., i]
+                        for i in 1:size(y)[end]]),)
     end
 end
 
 @adjoint function DiffEqArray(u, t)
     DiffEqArray(u, t),
     y -> begin
         y isa Ref && (y = VectorOfArray(y[].u))
-        (DiffEqArray([y[ntuple(x -> Colon(), ndims(y.u) - 1)..., i]
-                       for i in 1:size(y.u)[end]],
+        (DiffEqArray([y[ntuple(x -> Colon(), ndims(y) - 1)..., i]
+                       for i in 1:size(y)[end]],
             t), nothing)
     end
 end

diff --git a/src/RecursiveArrayTools.jl b/src/RecursiveArrayTools.jl
@@ -28,6 +28,7 @@ end
 
 import GPUArraysCore
 Base.convert(T::Type{<:GPUArraysCore.AnyGPUArray}, VA::AbstractVectorOfArray) = stack(VA.u)
+(T::Type{<:GPUArraysCore.AnyGPUArray})(VA::AbstractVectorOfArray) = T(Array(VA))
 
 import Requires
 @static if !isdefined(Base, :get_extension)