Remove unnecessary jax.vmap in model.link_bias_accelerations

ami-iit · Nov 21, 2024 · bf02ea2 · bf02ea2
1 parent 44daf21
commit bf02ea2
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/src/jaxsim/api/model.py b/src/jaxsim/api/model.py
@@ -1951,11 +1951,9 @@ def body_to_other_representation(
 
         case VelRepr.Mixed:
             W_H_L = js.model.forward_kinematics(model=model, data=data)
-            LW_H_L = jax.vmap(lambda W_H_L: W_H_L.at[0:3, 3].set(jnp.zeros(3)))(W_H_L)
+            LW_H_L = W_H_L.at[:, 0:3, 3].set(jnp.zeros(3))
             C_H_L = LW_H_L
-            L_v_CL = L_v_LW_L = jax.vmap(  # noqa: F841
-                lambda v: v.at[0:3].set(jnp.zeros(3))
-            )(L_v_WL)
+            L_v_CL = L_v_LW_L = L_v_WL.at[:, 0:3].set(jnp.zeros(3))  # noqa: F841
 
         case _:
             raise ValueError(data.velocity_representation)