fix for NONCONST_CLS

AMICI-dev · Oct 24, 2024 · c548c93 · c548c93
1 parent 7c3aef9
commit c548c93
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/python/sdist/amici/jax.py b/python/sdist/amici/jax.py
@@ -167,7 +167,9 @@ def _run(
                 ts, ps, k, x0, checkpointed=checkpointed
             )
         else:
-            x = tuple(jnp.array([x0_i] * len(ts)) for x0_i in x0)
+            x = tuple(
+                self.x_solver(jnp.array([x0_i] * len(ts)) for x0_i in x0)
+            )
             tcl = self.tcl(x0, ps, k)
             stats = None
         obs = self._obs(ts, x, ps, k, tcl)