Remove JAX deprecated methods

ami-iit · Nov 14, 2024 · f00c3d4 · f00c3d4
1 parent b3c59ea
commit f00c3d4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/jaxsim/integrators/common.py b/src/jaxsim/integrators/common.py
@@ -319,7 +319,7 @@ def _compute_next_state(
         f = lambda x, t: self.dynamics(x=x, t=t, **kwargs)
 
         # Initialize the carry of the for loop with the stacked kᵢ vectors.
-        carry0 = jax.tree_map(
+        carry0 = jax.tree.map(
             lambda l: jnp.zeros((c.size, *l.shape), dtype=l.dtype), x0
         )