Fix ReverseDiff downstream

SciML · May 14, 2024 · c868913 · c868913
1 parent 803d02e
commit c868913
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "RecursiveArrayTools"
 uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "3.17.0"
+version = "3.18.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/ext/RecursiveArrayToolsReverseDiffExt.jl b/ext/RecursiveArrayToolsReverseDiffExt.jl
@@ -3,12 +3,13 @@ module RecursiveArrayToolsReverseDiffExt
 using RecursiveArrayTools
 using ReverseDiff
 using Zygote: @adjoint
+using RecursiveArrayTools.ArrayInterface
 
 function trackedarraycopyto!(dest, src)
     for (i, slice) in zip(eachindex(dest.u), eachslice(src, dims = ndims(src)))
         if dest.u[i] isa AbstractArray
-            dest.u[i] = reshape(reduce(vcat, slice), size(dest.u[i]))
-        else
+            dest.u[i] = reshape(ArrayInterface.aos_to_soa(slice), size(dest.u[i]))
+        elseif dest.u[i]
             trackedarraycopyto!(dest.u[i], slice)
         end
     end