Merge pull request #808 from AayushSabharwal/as/fix-tests

fix: use `split = false` system for remake autodiff tests
SciML · Oct 7, 2024 · 70ae7fb · 70ae7fb
2 parents b064f2b + 6184c69
commit 70ae7fb
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/ext/SciMLBaseZygoteExt.jl b/ext/SciMLBaseZygoteExt.jl
@@ -106,11 +106,12 @@ end
         if is_observed(VA, sym)
             f = observed(VA, sym)
             p = parameter_values(VA)
-            tunables, _, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
+            tunables, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
             u = state_values(VA)
             t = current_time(VA)
             y, back = Zygote.pullback(u, tunables) do u, tunables
-                f.(u, Ref(tunables), t)
+                _p = repack(tunables)
+                f.(u, Ref(_p), t)
             end
             gs = back(Δ)
             (gs[1], nothing)

diff --git a/test/downstream/remake_autodiff.jl b/test/downstream/remake_autodiff.jl
@@ -14,11 +14,11 @@ function lotka_volterra(; name = name)
 end
 
 @named lotka_volterra_sys = lotka_volterra()
-lotka_volterra_sys = structural_simplify(lotka_volterra_sys)
+lotka_volterra_sys = structural_simplify(lotka_volterra_sys, split = false)
 prob = ODEProblem(lotka_volterra_sys, [], (0.0, 10.0), [])
 sol = solve(prob, Tsit5(), reltol = 1e-6, abstol = 1e-6)
-u0 = [1.0 1.0]
-p = [1.5 1.0 1.0 1.0]
+u0 = [1.0, 1.0]
+p = [1.5, 1.0, 1.0, 1.0]
 
 function sum_of_solution(u0, p)
     _prob = remake(prob, u0 = u0, p = p)