Address deprecation warning

ami-iit · Jun 19, 2024 · 7250666 · 7250666
1 parent 2dc884f
commit 7250666
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/jaxsim/integrators/common.py b/src/jaxsim/integrators/common.py
@@ -422,7 +422,9 @@ def compute_ki() -> jax.Array:
 
         # Update the FSAL property for the next iteration.
         if self.has_fsal:
-            self.params["dxdt0"] = jax.tree_map(lambda l: l[self.index_of_fsal], K)
+            self.params["dxdt0"] = jax.tree_util.tree_map(
+                lambda l: l[self.index_of_fsal], K
+            )
 
         # Compute the output state.
         # Note that z contains as many new states as the rows of `b.T`.