Update SciMLBaseZygoteExt.jl

SciML · Oct 28, 2023 · 30c0342 · 30c0342
1 parent 9cb8c14
commit 30c0342
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/ext/SciMLBaseZygoteExt.jl b/ext/SciMLBaseZygoteExt.jl
@@ -56,18 +56,18 @@ end
     VA[sym, j], ODESolution_getindex_pullback
 end
 
-ZygoteRules.@adjoint function DiffEqBase.EnsembleSolution(sim, time, converged)
+ZygoteRules.@adjoint function DiffEqBase.EnsembleSolution(sim, time, converged, stats)
     out = EnsembleSolution(sim, time, converged)
     function EnsembleSolution_adjoint(p̄::AbstractArray{T, N}) where {T, N}
         arrarr = [[p̄[ntuple(x -> Colon(), Val(N - 2))..., j, i]
                    for j in 1:size(p̄)[end - 1]] for i in 1:size(p̄)[end]]
-        (EnsembleSolution(arrarr, 0.0, true), nothing, nothing)
+        (EnsembleSolution(arrarr, 0.0, true), nothing, nothing, nothing)
     end
     function EnsembleSolution_adjoint(p̄::AbstractArray{<:AbstractArray, 1})
-        (EnsembleSolution(p̄, 0.0, true), nothing, nothing)
+        (EnsembleSolution(p̄, 0.0, true), nothing, nothing, nothing)
     end
     function EnsembleSolution_adjoint(p̄::EnsembleSolution)
-        (p̄, nothing, nothing)
+        (p̄, nothing, nothing, nothing)
     end
     out, EnsembleSolution_adjoint
 end