Fix Array of Structs to Struct of Array in ReverseDiff

Fixes #581
SciML · Mar 12, 2024 · a6737f7 · a6737f7
1 parent 4a86cef
commit a6737f7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/concrete_solve.jl b/src/concrete_solve.jl
@@ -1313,9 +1313,9 @@ function DiffEqBase._concrete_solve_adjoint(
             end
         else
             # use TrackedArray for efficiency of the tape
-            _f(args...) = reduce(vcat, prob.f(args...))
+            _f(args...) = ArrayInterface.aos_to_soa(prob.f(args...))
             if prob isa SDEProblem
-                _g(args...) = reduce(vcat, prob.g(args...))
+                _g(args...) = ArrayInterface.aos_to_soa(prob.g(args...))
                 _prob = remake(prob,
                     f = DiffEqBase.parameterless_type(prob.f){
                         SciMLBase.isinplace(prob),