jaxsim 0.4.3.dev280__py3-none-any.whl → 0.4.3.dev282__py3-none-any.whl

jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.4.3.dev280'
-__version_tuple__ = version_tuple = (0, 4, 3, 'dev280')
+__version__ = version = '0.4.3.dev282'
+__version_tuple__ = version_tuple = (0, 4, 3, 'dev282')

jaxsim/math/rotation.py

@@ -1,4 +1,3 @@
-import jax
 import jax.numpy as jnp
 import jaxlie
 
@@ -8,6 +7,7 @@ from .skew import Skew
 
 
 class Rotation:
+
     @staticmethod
     def x(theta: jtp.Float) -> jtp.Matrix:
         """
@@ -19,6 +19,7 @@ class Rotation:
         Returns:
             jtp.Matrix: 3D rotation matrix.
         """
+
         return jaxlie.SO3.from_x_radians(theta=theta).as_matrix()
 
     @staticmethod
@@ -32,6 +33,7 @@ class Rotation:
         Returns:
             jtp.Matrix: 3D rotation matrix.
         """
+
         return jaxlie.SO3.from_y_radians(theta=theta).as_matrix()
 
     @staticmethod
@@ -45,6 +47,7 @@ class Rotation:
         Returns:
             jtp.Matrix: 3D rotation matrix.
         """
+
         return jaxlie.SO3.from_z_radians(theta=theta).as_matrix()
 
     @staticmethod
@@ -53,17 +56,18 @@ class Rotation:
         Generate a 3D rotation matrix from an axis-angle representation.
 
         Args:
-            vector (jtp.Vector): Axis-angle representation as a 3D vector.
+            vector: Axis-angle representation or the rotation as a 3D vector.
 
         Returns:
-            jtp.Matrix: 3D rotation matrix.
-
+            The SO(3) rotation matrix.
         """
+
         vector = vector.squeeze()
-        theta = jnp.linalg.norm(vector)
 
-        def theta_is_not_zero(theta_and_v: tuple[jtp.Float, jtp.Vector]) -> jtp.Matrix:
-            theta, v = theta_and_v
+        def theta_is_not_zero(axis: jtp.Vector) -> jtp.Matrix:
+
+            v = axis
+            theta = jnp.linalg.norm(v)
 
             s = jnp.sin(theta)
             c = jnp.cos(theta)
@@ -77,9 +81,19 @@ class Rotation:
 
             return R.transpose()
 
-        return jax.lax.cond(
-            pred=(theta == 0.0),
-            true_fun=lambda operand: jnp.eye(3),
-            false_fun=theta_is_not_zero,
-            operand=(theta, vector),
+        # Use the double-where trick to prevent JAX problems when the
+        # jax.jit and jax.grad transforms are applied.
+        return jnp.where(
+            jnp.linalg.norm(vector) > 0,
+            theta_is_not_zero(
+                axis=jnp.where(
+                    jnp.linalg.norm(vector) > 0,
+                    vector,
+                    # The following line is a workaround to prevent division by 0.
+                    # Considering the outer where, this branch is never executed.
+                    jnp.ones(3),
+                )
+            ),
+            # Return an identity rotation matrix when the input vector is zero.
+            jnp.eye(3),
         )

{jaxsim-0.4.3.dev280.dist-info → jaxsim-0.4.3.dev282.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.4.3.dev280
+Version: 0.4.3.dev282
 Summary: A differentiable physics engine and multibody dynamics library for control and robot learning.
 Author-email: Diego Ferigo <dgferigo@gmail.com>
 Maintainer-email: Diego Ferigo <dgferigo@gmail.com>, Filippo Luca Ferretti <filippo.ferretti@iit.it>

{jaxsim-0.4.3.dev280.dist-info → jaxsim-0.4.3.dev282.dist-info}/RECORD

@@ -1,5 +1,5 @@
 jaxsim/__init__.py,sha256=opgtbhhd1kDsHI4H1vOd3loMPDRi884yQ3tohfFGfNc,3382
-jaxsim/_version.py,sha256=U12VcDXh7Zu7La30KnfhDG0g-i5JqVyuLLRlI8iTnrI,428
+jaxsim/_version.py,sha256=8Lzzuj01bYiR6tHmTBRLQW2sSRCWlBO6B-razI_HdjU,428
 jaxsim/exceptions.py,sha256=vSoScaRD4nvh6jltgK9Ry5pKnE0O5hb4_yI_pk_fvR8,2175
 jaxsim/logging.py,sha256=STI-D_upXZYX-ZezLrlJJ0UlD5YspST0vZ_DcIwkzO4,1553
 jaxsim/typing.py,sha256=2HXy9hgazPXjofi1vLQ09ZubPtgVmg80U9NKmZ6NYiI,761
@@ -26,7 +26,7 @@ jaxsim/math/cross.py,sha256=U7yEx_l75mSy5g6O-jsjBztApvxC3WaV4MpkS5tThu4,1330
 jaxsim/math/inertia.py,sha256=01hz6wMFreN2jBA0rVoBS1YMVh77KvwuzXSOpI3pxNk,1614
 jaxsim/math/joint_model.py,sha256=EzAveaG5B6ZnCFNUzN30KEQUVesd83lfWXJarYR-kUw,9989
 jaxsim/math/quaternion.py,sha256=_WA7W3iv7px83sWO1V1n0-J78hqAlO4SL1-jofE-UZ4,4754
-jaxsim/math/rotation.py,sha256=k-nwT79zmWrys3NNAB-lGWxat7Kqm_6JnFRoimJ8rBg,2156
+jaxsim/math/rotation.py,sha256=P34sx28Rh1MhNwxUxqXjxP-ZN1_5tvoflMoAIpy2LbE,2586
 jaxsim/math/skew.py,sha256=oOGSSR8PUGROl6IJFlrmu6K3gPH-u16hUPfKIkcVv9o,1177
 jaxsim/math/transform.py,sha256=KXzQgOnCfAtbXCwxhplpJ3F0JT3oEyeLVby1_uRAryQ,2892
 jaxsim/mujoco/__init__.py,sha256=fZyRWre49pIhOrYdf6yJk_hOax8qWGe8OCmoq-dMVq8,201
@@ -65,8 +65,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=TGmTQV2Lq7Q-2nLoAEaeNtkPa_qj0IKkdBm4COj46Os,11312
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=Fh82ZcaFi5fUnByyFLnmumaobsu1hJIvFdopUVzJ1ps,4052
-jaxsim-0.4.3.dev280.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.4.3.dev280.dist-info/METADATA,sha256=3Ad6_an0aEBTc9JVFxZbdbT-0PnAWLKgEVRW37ztcxo,17276
-jaxsim-0.4.3.dev280.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-jaxsim-0.4.3.dev280.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.4.3.dev280.dist-info/RECORD,,
+jaxsim-0.4.3.dev282.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.4.3.dev282.dist-info/METADATA,sha256=Llp_WIbsROSnrBpSMJUdFnAOOkLDobo5f87CO8xvOuA,17276
+jaxsim-0.4.3.dev282.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+jaxsim-0.4.3.dev282.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.4.3.dev282.dist-info/RECORD,,