fix: ChainRulesCore.ProjectTo implementation for VoA

SciML · Jan 5, 2024 · 8e5f902 · 8e5f902
1 parent 433a780
commit 8e5f902
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ext/RecursiveArrayToolsZygoteExt.jl b/ext/RecursiveArrayToolsZygoteExt.jl
@@ -125,7 +125,7 @@ end
 
 ChainRulesCore.ProjectTo(a::AbstractVectorOfArray) = ChainRulesCore.ProjectTo{VectorOfArray}((sz = size(a)))
 
-function (p::ChainRulesCore.ProjectTo{VectorOfArray})(x)
+function (p::ChainRulesCore.ProjectTo{VectorOfArray})(x::Union{AbstractArray,AbstractVectorOfArray})
     arr = reshape(x, p.sz)
     return VectorOfArray([arr[:, i] for i in 1:p.sz[end]])
 end