jaxsim 0.7.1.dev40__py3-none-any.whl → 0.7.1.dev46__py3-none-any.whl

jaxsim/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.7.1.dev40'
-__version_tuple__ = version_tuple = (0, 7, 1, 'dev40')
+__version__ = version = '0.7.1.dev46'
+__version_tuple__ = version_tuple = (0, 7, 1, 'dev46')

jaxsim/math/utils.py

@@ -1,8 +1,47 @@
+import jax
 import jax.numpy as jnp
 
 import jaxsim.typing as jtp
 
 
+def _make_safe_norm(axis, keepdims):
+    @jax.custom_jvp
+    def _safe_norm(array: jtp.ArrayLike) -> jtp.Array:
+        """
+        Compute an array norm handling NaNs and making sure that
+        it is safe to get the gradient.
+
+        Args:
+            array: The array for which to compute the norm.
+
+        Returns:
+            The norm of the array with handling for zero arrays to avoid NaNs.
+        """
+        # Compute the norm of the array along the specified axis.
+        return jnp.linalg.norm(array, axis=axis, keepdims=keepdims)
+
+    @_safe_norm.defjvp
+    def _safe_norm_jvp(primals, tangents):
+        (x,), (x_dot,) = primals, tangents
+
+        # Check if the entire array is composed of zeros.
+        is_zero = jnp.all(x == 0.0)
+
+        # Replace zeros with an array of ones temporarily to avoid division by zero.
+        # This ensures the computation of norm does not produce NaNs or Infs.
+        array = jnp.where(is_zero, jnp.ones_like(x), x)
+
+        # Compute the norm of the array along the specified axis.
+        norm = jnp.linalg.norm(array, axis=axis, keepdims=keepdims)
+
+        dot = jnp.sum(array * x_dot, axis=axis, keepdims=keepdims)
+        tangent = jnp.where(is_zero, 0.0, dot / norm)
+
+        return jnp.where(is_zero, 0.0, norm), tangent
+
+    return _safe_norm
+
+
 def safe_norm(array: jtp.ArrayLike, *, axis=None, keepdims: bool = False) -> jtp.Array:
     """
     Compute an array norm handling NaNs and making sure that
@@ -16,17 +55,4 @@ def safe_norm(array: jtp.ArrayLike, *, axis=None, keepdims: bool = False) -> jtp
     Returns:
         The norm of the array with handling for zero arrays to avoid NaNs.
     """
-
-    # Check if the entire array is composed of zeros.
-    is_zero = jnp.allclose(array, 0.0)
-
-    # Replace zeros with an array of ones temporarily to avoid division by zero.
-    # This ensures the computation of norm does not produce NaNs or Infs.
-    array = jnp.where(is_zero, jnp.ones_like(array), array)
-
-    # Compute the norm of the array along the specified axis.
-    norm = jnp.linalg.norm(array, axis=axis, keepdims=keepdims)
-
-    # Use `jnp.where` to set the norm to 0.0 where the input array was all zeros.
-    # This usage supports potential batch processing for future scalability.
-    return jnp.where(is_zero, 0.0, norm)
+    return _make_safe_norm(axis, keepdims)(array)

{jaxsim-0.7.1.dev40.dist-info → jaxsim-0.7.1.dev46.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxsim
-Version: 0.7.1.dev40
+Version: 0.7.1.dev46
 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.7.1.dev40.dist-info → jaxsim-0.7.1.dev46.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=EKeysKN-7UswwJLCl7n6qIBKQIVUtYsCMYu_tCoFn7g,3628
-jaxsim/_version.py,sha256=OwTc5XXxttxZvtIDUc-qIEIXFTKMLrNesK0EXIP8otg,526
+jaxsim/_version.py,sha256=zEa993Q8cwGHCYf324xAXfDMwJV-C60H3PCcfC9TMNI,526
 jaxsim/exceptions.py,sha256=MQ3LRMfVMX2-g3qYj7mUVNV9OLlIA48TANJegbcQyXI,2641
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=7msl8t5Jt09RNYfKdPJtpjLfWurldcycDappb045Eso,761
@@ -26,7 +26,7 @@ jaxsim/math/quaternion.py,sha256=MSaZywzJDxs2te1ZELeIcupKSFIA9q_pdXy7fDAEqM4,453
 jaxsim/math/rotation.py,sha256=TEUtT3X2tFieNxdlccup1pfaTgCTtfX-hTNotd8-nNk,1892
 jaxsim/math/skew.py,sha256=z_9YN-NDHL3n4KXWNbzTSMkFDZ0SDpz4RUcwwYFOaao,1402
 jaxsim/math/transform.py,sha256=d0_m_obmUOmnI8Bte0ktvibR9Hv9M9qpg8tVuLON2g0,3192
-jaxsim/math/utils.py,sha256=IiH01iN54BtLnULC04pDfYe8Av99p3FGdYp2jJInm30,1166
+jaxsim/math/utils.py,sha256=JgJrBPeuCvi0969VqoNsyk3CflQiLzopngKDjl6RfiE,1898
 jaxsim/mujoco/__init__.py,sha256=1kAWzYOS7nP29S5FGyWPqiAnPf4yPSoaPW-WBGBjVV0,214
 jaxsim/mujoco/__main__.py,sha256=GBmB7J-zj75ZnFyuAAmpSOpbxi_HhHhWJeot3ljGDJY,5291
 jaxsim/mujoco/loaders.py,sha256=OCk1T11iIm3qZUibNpo_bxxLgaGSkCpLt7ae_ND0ExA,23272
@@ -65,8 +65,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=XzmZeIibcaOzaxpprsGSxH3UrM66PAO456rFV91sNXg,11453
 jaxsim/utils/tracing.py,sha256=Btwxdfhb7fJLk3r5PlQkGYj60Y2KbFT1gANGIA697FU,530
 jaxsim/utils/wrappers.py,sha256=3IMwydqFgmSPqeuUQ3PRmdhDc1IoT6XC23jPC_LjWXs,4175
-jaxsim-0.7.1.dev40.dist-info/licenses/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.7.1.dev40.dist-info/METADATA,sha256=UgEDJHiMs785VbvpiOcHA4uAKb0zlRYFa7GSvpqVadY,17851
-jaxsim-0.7.1.dev40.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-jaxsim-0.7.1.dev40.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.7.1.dev40.dist-info/RECORD,,
+jaxsim-0.7.1.dev46.dist-info/licenses/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.7.1.dev46.dist-info/METADATA,sha256=BC1OvFrpPoiO0y1te_UYwd4IMoZod5i9HWmsKOBZqAs,17851
+jaxsim-0.7.1.dev46.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
+jaxsim-0.7.1.dev46.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.7.1.dev46.dist-info/RECORD,,

{jaxsim-0.7.1.dev40.dist-info → jaxsim-0.7.1.dev46.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.7.1)
+Generator: setuptools (80.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any