fix: use split = false in de.jit

SciML · Nov 20, 2024 · b7d7391 · b7d7391
1 parent 2ccb76d
commit b7d7391
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/diffeqpy/de.py b/diffeqpy/de.py
@@ -2,6 +2,11 @@
 from . import load_julia_packages
 de, _, _ = load_julia_packages("DifferentialEquations", "ModelingToolkit", "PythonCall")
 from juliacall import Main
-de.jit = Main.seval("jit(x) = typeof(x).name.wrapper(ModelingToolkit.complete(ModelingToolkit.modelingtoolkitize(x)), x.u0, x.tspan, x.p)") # kinda hackey
-de.jit32 = Main.seval("jit(x) = typeof(x).name.wrapper(ModelingToolkit.complete(ModelingToolkit.modelingtoolkitize(x)), Float32.(x.u0), Float32.(x.tspan), Float32.(x.p))") # kinda hackey
+de.jit = Main.seval("jit(x) = typeof(x).name.wrapper(ModelingToolkit.complete(ModelingToolkit.modelingtoolkitize(x); split = false), [], x.tspan)") # kinda hackey
+de.jit32 = Main.seval("""
+                      function jit(x)
+                          prob = typeof(x).name.wrapper(ModelingToolkit.complete(ModelingToolkit.modelingtoolkitize(x); split = false), [], Float32.(x.tspan))
+                          remake(prob; u0 = Float32.(prob.u0), p = Float32.(prob.p))
+                      end
+                      """) # kinda hackey
 sys.modules[__name__] = de # mutate myself