jaxsim 0.5.1.dev56__tar.gz → 0.5.1.dev60__tar.gz

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev56
+Version: 0.5.1.dev60
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.5.1.dev56'
-__version_tuple__ = version_tuple = (0, 5, 1, 'dev56')
+__version__ = version = '0.5.1.dev60'
+__version_tuple__ = version_tuple = (0, 5, 1, 'dev60')

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim/api/com.py

@@ -8,6 +8,7 @@ import jaxsim.typing as jtp
 from .common import VelRepr
 
 
+@js.common.named_scope
 @jax.jit
 def com_position(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -44,6 +45,7 @@ def com_position(
     return (W_H_B @ B_p̃_CoM)[0:3].astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def com_linear_velocity(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -74,6 +76,7 @@ def com_linear_velocity(
     return G_vl_WG
 
 
+@js.common.named_scope
 @jax.jit
 def centroidal_momentum(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -101,6 +104,7 @@ def centroidal_momentum(
     return G_J @ ν
 
 
+@js.common.named_scope
 @jax.jit
 def centroidal_momentum_jacobian(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -149,6 +153,7 @@ def centroidal_momentum_jacobian(
     return G_Xf_B @ B_Jh
 
 
+@js.common.named_scope
 @jax.jit
 def locked_centroidal_spatial_inertia(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -186,6 +191,7 @@ def locked_centroidal_spatial_inertia(
     return G_Xf_B @ B_Mbb_B @ B_Xv_G
 
 
+@js.common.named_scope
 @jax.jit
 def average_centroidal_velocity(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -213,6 +219,7 @@ def average_centroidal_velocity(
     return G_J @ ν
 
 
+@js.common.named_scope
 @jax.jit
 def average_centroidal_velocity_jacobian(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -239,6 +246,7 @@ def average_centroidal_velocity_jacobian(
     return jnp.linalg.inv(G_Mbb) @ G_J
 
 
+@js.common.named_scope
 @jax.jit
 def bias_acceleration(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim/api/common.py

@@ -3,7 +3,8 @@ import contextlib
 import dataclasses
 import enum
 import functools
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
+from typing import ParamSpec, TypeVar
 
 import jax
 import jax.numpy as jnp
@@ -20,6 +21,21 @@ except ImportError:
     from typing_extensions import Self
 
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+def named_scope(fn, name: str | None = None) -> Callable[_P, _R]:
+    """Applies a JAX named scope to a function for improved profiling and clarity."""
+
+    @functools.wraps(fn)
+    def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R:
+        with jax.named_scope(name or fn.__name__):
+            return fn(*args, **kwargs)
+
+    return wrapper
+
+
 @enum.unique
 class VelRepr(enum.IntEnum):
     """

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim/api/contact.py

@@ -16,6 +16,7 @@ from jaxsim.rbda import contacts
 from .common import VelRepr
 
 
+@js.common.named_scope
 @jax.jit
 def collidable_point_kinematics(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -52,6 +53,7 @@ def collidable_point_kinematics(
     return W_p_Ci, W_ṗ_Ci
 
 
+@js.common.named_scope
 @jax.jit
 def collidable_point_positions(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -72,6 +74,7 @@ def collidable_point_positions(
     return W_p_Ci
 
 
+@js.common.named_scope
 @jax.jit
 def collidable_point_velocities(
     model: js.model.JaxSimModel, data: js.data.JaxSimModelData
@@ -92,6 +95,7 @@ def collidable_point_velocities(
     return W_ṗ_Ci
 
 
+@js.common.named_scope
 @jax.jit
 def collidable_point_forces(
     model: js.model.JaxSimModel,
@@ -129,6 +133,7 @@ def collidable_point_forces(
     return f_Ci
 
 
+@js.common.named_scope
 @jax.jit
 def collidable_point_dynamics(
     model: js.model.JaxSimModel,
@@ -227,6 +232,7 @@ def collidable_point_dynamics(
     return f_Ci, aux_data
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["link_names"])
 def in_contact(
     model: js.model.JaxSimModel,
@@ -424,6 +430,7 @@ def estimate_good_contact_parameters(
     return parameters
 
 
+@js.common.named_scope
 @jax.jit
 def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Array:
     r"""
@@ -469,6 +476,7 @@ def transforms(model: js.model.JaxSimModel, data: js.data.JaxSimModelData) -> jt
     return jax.vmap(lambda W_H_Li, L_H_Ci: W_H_Li @ L_H_Ci)(W_H_L, L_H_C)
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian(
     model: js.model.JaxSimModel,
@@ -561,6 +569,7 @@ def jacobian(
     return O_J_WC
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian_derivative(
     model: js.model.JaxSimModel,

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim/api/data.py

@@ -253,6 +253,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return -self.gravity[2]
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def joint_positions(
         self,
@@ -300,6 +301,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.state.physics_model.joint_positions[joint_idxs]
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def joint_velocities(
         self,
@@ -347,6 +349,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.state.physics_model.joint_velocities[joint_idxs]
 
+    @js.common.named_scope
     @jax.jit
     def base_position(self) -> jtp.Vector:
         """
@@ -358,6 +361,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.state.physics_model.base_position.squeeze()
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["dcm"])
     def base_orientation(self, dcm: jtp.BoolLike = False) -> jtp.Vector | jtp.Matrix:
         """
@@ -386,6 +390,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             float
         )
 
+    @js.common.named_scope
     @jax.jit
     def base_transform(self) -> jtp.Matrix:
         """
@@ -405,6 +410,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ]
         )
 
+    @js.common.named_scope
     @jax.jit
     def base_velocity(self) -> jtp.Vector:
         """
@@ -434,6 +440,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             .astype(float)
         )
 
+    @js.common.named_scope
     @jax.jit
     def generalized_position(self) -> tuple[jtp.Matrix, jtp.Vector]:
         r"""
@@ -446,6 +453,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
 
         return self.base_transform(), self.joint_positions()
 
+    @js.common.named_scope
     @jax.jit
     def generalized_velocity(self) -> jtp.Vector:
         r"""
@@ -466,6 +474,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
     # Store quantities
     # ================
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def reset_joint_positions(
         self,
@@ -514,6 +523,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             s=self.state.physics_model.joint_positions.at[joint_idxs].set(positions)
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def reset_joint_velocities(
         self,
@@ -562,6 +572,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ṡ=self.state.physics_model.joint_velocities.at[joint_idxs].set(velocities)
         )
 
+    @js.common.named_scope
     @jax.jit
     def reset_base_position(self, base_position: jtp.VectorLike) -> Self:
         """
@@ -585,6 +596,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ),
         )
 
+    @js.common.named_scope
     @jax.jit
     def reset_base_quaternion(self, base_quaternion: jtp.VectorLike) -> Self:
         """
@@ -612,6 +624,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             ),
         )
 
+    @js.common.named_scope
     @jax.jit
     def reset_base_pose(self, base_pose: jtp.MatrixLike) -> Self:
         """
@@ -634,6 +647,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             base_quaternion=W_Q_B
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["velocity_representation"])
     def reset_base_linear_velocity(
         self,
@@ -665,6 +679,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             velocity_representation=velocity_representation,
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["velocity_representation"])
     def reset_base_angular_velocity(
         self,
@@ -696,6 +711,7 @@ class JaxSimModelData(common.ModelDataWithVelocityRepresentation):
             velocity_representation=velocity_representation,
         )
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["velocity_representation"])
     def reset_base_velocity(
         self,

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim/api/frame.py

@@ -16,6 +16,7 @@ from .common import VelRepr
 # =======================
 
 
+@js.common.named_scope
 @jax.jit
 def idx_of_parent_link(
     model: js.model.JaxSimModel, *, frame_index: jtp.IntLike
@@ -45,6 +46,7 @@ def idx_of_parent_link(
     ]
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames="frame_name")
 def name_to_idx(model: js.model.JaxSimModel, *, frame_name: str) -> jtp.Int:
     """
@@ -97,6 +99,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, frame_index: jtp.IntLike) -> str
     ]
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["frame_names"])
 def names_to_idxs(
     model: js.model.JaxSimModel, *, frame_names: Sequence[str]
@@ -139,6 +142,7 @@ def idxs_to_names(
 # ==========
 
 
+@js.common.named_scope
 @jax.jit
 def transform(
     model: js.model.JaxSimModel,
@@ -180,6 +184,7 @@ def transform(
     return W_H_L @ L_H_F
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def velocity(
     model: js.model.JaxSimModel,
@@ -230,6 +235,7 @@ def velocity(
     return O_J_WF_I @ I_ν
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian(
     model: js.model.JaxSimModel,
@@ -309,6 +315,7 @@ def jacobian(
     return O_J_WL_I
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["output_vel_repr"])
 def jacobian_derivative(
     model: js.model.JaxSimModel,

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim/api/joint.py

@@ -13,6 +13,7 @@ from jaxsim import exceptions
 # =======================
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames="joint_name")
 def name_to_idx(model: js.model.JaxSimModel, *, joint_name: str) -> jtp.Int:
     """
@@ -61,6 +62,7 @@ def idx_to_name(model: js.model.JaxSimModel, *, joint_index: jtp.IntLike) -> str
     return model.kin_dyn_parameters.joint_model.joint_names[joint_index + 1]
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames="joint_names")
 def names_to_idxs(
     model: js.model.JaxSimModel, *, joint_names: Sequence[str]
@@ -141,6 +143,7 @@ def position_limit(
     return s_min.astype(float), s_max.astype(float)
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["joint_names"])
 def position_limits(
     model: js.model.JaxSimModel, *, joint_names: Sequence[str] | None = None
@@ -176,6 +179,7 @@ def position_limits(
 # ======================
 
 
+@js.common.named_scope
 @functools.partial(jax.jit, static_argnames=["joint_names"])
 def random_joint_positions(
     model: js.model.JaxSimModel,

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim/api/model.py

@@ -491,6 +491,7 @@ def reduce(
 # ===================
 
 
+@js.common.named_scope
 @jax.jit
 def total_mass(model: JaxSimModel) -> jtp.Float:
     """
@@ -506,6 +507,7 @@ def total_mass(model: JaxSimModel) -> jtp.Float:
     return model.kin_dyn_parameters.link_parameters.mass.sum().astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def link_spatial_inertia_matrices(model: JaxSimModel) -> jtp.Array:
     """
@@ -528,6 +530,7 @@ def link_spatial_inertia_matrices(model: JaxSimModel) -> jtp.Array:
 # ==============================
 
 
+@js.common.named_scope
 @jax.jit
 def forward_kinematics(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Array:
     """
@@ -915,6 +918,7 @@ def forward_dynamics(
     )
 
 
+@js.common.named_scope
 @jax.jit
 def forward_dynamics_aba(
     model: JaxSimModel,
@@ -1059,6 +1063,7 @@ def forward_dynamics_aba(
     return C_v̇_WB.astype(float), s̈.astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def forward_dynamics_crb(
     model: JaxSimModel,
@@ -1150,6 +1155,7 @@ def forward_dynamics_crb(
     return v̇_WB, s̈
 
 
+@js.common.named_scope
 @jax.jit
 def free_floating_mass_matrix(
     model: JaxSimModel, data: js.data.JaxSimModelData
@@ -1195,6 +1201,7 @@ def free_floating_mass_matrix(
             raise ValueError(data.velocity_representation)
 
 
+@js.common.named_scope
 @jax.jit
 def free_floating_coriolis_matrix(
     model: JaxSimModel, data: js.data.JaxSimModelData
@@ -1311,6 +1318,7 @@ def free_floating_coriolis_matrix(
             raise ValueError(data.velocity_representation)
 
 
+@js.common.named_scope
 @jax.jit
 def inverse_dynamics(
     model: JaxSimModel,
@@ -1466,6 +1474,7 @@ def inverse_dynamics(
     return f_B.astype(float), τ.astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def free_floating_gravity_forces(
     model: JaxSimModel, data: js.data.JaxSimModelData
@@ -1515,6 +1524,7 @@ def free_floating_gravity_forces(
     ).astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def free_floating_bias_forces(
     model: JaxSimModel, data: js.data.JaxSimModelData
@@ -1584,6 +1594,7 @@ def free_floating_bias_forces(
 # ==========================
 
 
+@js.common.named_scope
 @jax.jit
 def locked_spatial_inertia(
     model: JaxSimModel, data: js.data.JaxSimModelData
@@ -1602,6 +1613,7 @@ def locked_spatial_inertia(
     return total_momentum_jacobian(model=model, data=data)[:, 0:6]
 
 
+@js.common.named_scope
 @jax.jit
 def total_momentum(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vector:
     """
@@ -1690,6 +1702,7 @@ def total_momentum_jacobian(
             raise ValueError(output_vel_repr)
 
 
+@js.common.named_scope
 @jax.jit
 def average_velocity(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Vector:
     """
@@ -1778,6 +1791,7 @@ def average_velocity_jacobian(
 # ========================
 
 
+@js.common.named_scope
 @jax.jit
 def link_bias_accelerations(
     model: JaxSimModel,
@@ -1987,6 +2001,7 @@ def link_bias_accelerations(
     return O_v̇_WL
 
 
+@js.common.named_scope
 @jax.jit
 def link_contact_forces(
     model: js.model.JaxSimModel,
@@ -2062,6 +2077,7 @@ def link_contact_forces(
 # ======
 
 
+@js.common.named_scope
 @jax.jit
 def mechanical_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Float:
     """
@@ -2081,6 +2097,7 @@ def mechanical_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.
     return (K + U).astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def kinetic_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Float:
     """
@@ -2102,6 +2119,7 @@ def kinetic_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Flo
     return K.squeeze().astype(float)
 
 
+@js.common.named_scope
 @jax.jit
 def potential_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.Float:
     """
@@ -2128,6 +2146,7 @@ def potential_energy(model: JaxSimModel, data: js.data.JaxSimModelData) -> jtp.F
 # ==========
 
 
+@js.common.named_scope
 @jax.jit
 def step(
     model: JaxSimModel,

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim/api/ode.py

@@ -85,6 +85,7 @@ def wrap_system_dynamics_for_integration(
 # ==================================
 
 
+@js.common.named_scope
 @jax.jit
 def system_velocity_dynamics(
     model: js.model.JaxSimModel,
@@ -331,6 +332,7 @@ def system_acceleration(
     return v̇_WB, s̈
 
 
+@js.common.named_scope
 @jax.jit
 def system_position_dynamics(
     model: js.model.JaxSimModel,
@@ -370,6 +372,7 @@ def system_position_dynamics(
     return W_ṗ_B, W_Q̇_B, ṡ
 
 
+@js.common.named_scope
 @jax.jit
 def system_dynamics(
     model: js.model.JaxSimModel,

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim/api/references.py

@@ -163,7 +163,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
     # ==================
     # Extract quantities
     # ==================
-
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["link_names"])
     def link_forces(
         self,
@@ -301,7 +301,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
     # ================
     # Store quantities
     # ================
-
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["joint_names"])
     def set_joint_force_references(
         self,
@@ -346,6 +346,7 @@ class JaxSimModelReferences(js.common.ModelDataWithVelocityRepresentation):
 
         return replace(forces=self._joint_force_references.at[joint_idxs].set(forces))
 
+    @js.common.named_scope
     @functools.partial(jax.jit, static_argnames=["link_names", "additive"])
     def apply_link_forces(
         self,

{jaxsim-0.5.1.dev56 → jaxsim-0.5.1.dev60}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.5.1.dev56
+Version: 0.5.1.dev60
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippoluca.ferretti@outlook.com>
 Maintainer-email: Filippo Luca Ferretti <filippo.ferretti@iit.it>, Alessandro Croci <alessandro.croci@iit.it>