Fix Array of Structs to Struct of Array in ReverseDiff

Fixes #581
SciML · Mar 4, 2024 · c0e15c5 · c0e15c5
1 parent 7386d1c
commit c0e15c5
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/concrete_solve.jl b/src/concrete_solve.jl
@@ -1309,9 +1309,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),