fix gradients

AMICI-dev · Oct 25, 2024 · 480b75a · 480b75a
1 parent 5a86f4c
commit 480b75a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/python/sdist/amici/jax.py b/python/sdist/amici/jax.py
@@ -237,7 +237,7 @@ def srun(
         dynamic=True,
     ):
         (llh, (x, obs, stats)), sllh = (
-            jax.value_and_grad(self._run, 1, True)
+            jax.value_and_grad(self._run, 2, True)
         )(ts, ts_dyn, p, k, k_preeq, my, pscale, dynamic=dynamic)
         return llh, sllh, (x, obs, stats)
 
@@ -254,10 +254,10 @@ def s2run(
         dynamic=True,
     ):
         (llh, (x, obs, stats)), sllh = (
-            jax.value_and_grad(self._run, 1, True)
+            jax.value_and_grad(self._run, 2, True)
         )(ts, ts_dyn, p, k, k_preeq, my, pscale, dynamic=dynamic)
 
-        s2llh = jax.hessian(self._run, 1, True)(
+        s2llh = jax.hessian(self._run, 2, True)(
             ts, ts_dyn, p, k, k_preeq, my, pscale, dynamic=dynamic
         )