Make output_vel_repr non-static in api.model

ami-iit · May 23, 2024 · 9bc8e2c · 9bc8e2c
1 parent a262dda
commit 9bc8e2c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/jaxsim/api/model.py b/src/jaxsim/api/model.py
@@ -400,7 +400,7 @@ def forward_kinematics(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp
     return jnp.atleast_3d(W_H_LL).astype(float)
 
 
-@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+@jax.jit
 def generalized_free_floating_jacobian(
     model: JaxSimModel,
     data: js.data.JaxSimModelData,
@@ -1125,7 +1125,7 @@ def total_momentum(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vec
     return Jh @ ν
 
 
-@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+@jax.jit
 def total_momentum_jacobian(
     model: JaxSimModel,
     data: js.data.JaxSimModelData,
@@ -1212,7 +1212,7 @@ def average_velocity(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.V
     return J @ ν
 
 
-@functools.partial(jax.jit, static_argnames=["output_vel_repr"])
+@jax.jit
 def average_velocity_jacobian(
     model: JaxSimModel,
     data: js.data.JaxSimModelData,