jaxsim 0.7.1.dev40__tar.gz → 0.7.1.dev46__tar.gz

{jaxsim-0.7.1.dev40 → jaxsim-0.7.1.dev46}/.github/workflows/gpu_benchmark.yml

@@ -1,6 +1,9 @@
 name: GPU Benchmarks
 
 on:
+  push:
+    branches:
+      - main
   pull_request:
     types: [opened, reopened, synchronize]
   workflow_dispatch:
@@ -46,7 +49,7 @@ jobs:
         uses: actions/cache/restore@v4
         with:
           path: ./cache
-          key: ${{ steps.get-main-branch-sha.outputs.sha }}-${{ runner.os }}-benchmark
+          key: ${{ runner.os }}-benchmark
 
       - name: Ensure version file is written
         run: |
@@ -60,7 +63,6 @@ jobs:
 
       - name: Compare benchmark results with main branch
         uses: benchmark-action/github-action-benchmark@v1.20.4
-        if: steps.cache.outputs.cache-hit == 'true'
         with:
           tool: 'pytest'
           output-file-path: output.json
@@ -106,4 +108,4 @@ jobs:
         if: ${{ github.ref_name == 'main' }}
         with:
           path: ./cache
-          key: ${{ steps.get-main-branch-sha.outputs.sha }}-${{ runner.os }}-benchmark
+          key: ${{ runner.os }}-benchmark

{jaxsim-0.7.1.dev40 → jaxsim-0.7.1.dev46}/PKG-INFO

@@ -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 → jaxsim-0.7.1.dev46}/src/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-0.7.1.dev46/src/jaxsim/math/utils.py

@@ -0,0 +1,58 @@
+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
+    it is safe to get the gradient.
+
+    Args:
+        array: The array for which to compute the norm.
+        axis: The axis for which to compute the norm.
+        keepdims: Whether to keep the dimensions of the input
+
+    Returns:
+        The norm of the array with handling for zero arrays to avoid NaNs.
+    """
+    return _make_safe_norm(axis, keepdims)(array)

{jaxsim-0.7.1.dev40 → jaxsim-0.7.1.dev46}/src/jaxsim.egg-info/PKG-INFO

@@ -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/src/jaxsim/math/utils.py

@@ -1,32 +0,0 @@
-import jax.numpy as jnp
-
-import jaxsim.typing as jtp
-
-
-def safe_norm(array: jtp.ArrayLike, *, axis=None, keepdims: bool = False) -> 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.
-        axis: The axis for which to compute the norm.
-        keepdims: Whether to keep the dimensions of the input
-
-    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)