fix state ids

AMICI-dev · Dec 5, 2024 · 757ffa1 · 757ffa1
1 parent c6affff
commit 757ffa1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/sdist/amici/jax/ode_export.py b/python/sdist/amici/jax/ode_export.py
@@ -226,7 +226,7 @@ def _generate_jax_code(self) -> None:
             # assign named variables from a jax array
             **_jax_variable_assignments(self.model, sym_names),
             # tuple of variable names (ids as they are unique)
-            **_jax_variable_ids(self.model, ("p", "k", "y", "x")),
+            **_jax_variable_ids(self.model, ("p", "k", "y", "x_rdata")),
             **{
                 "MODEL_NAME": self.model_name,
                 # keep track of the API version that the model was generated with so we