jaxsim 0.6.1.dev9__tar.gz → 0.6.1.dev13__tar.gz

{jaxsim-0.6.1.dev9 → jaxsim-0.6.1.dev13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jaxsim
-Version: 0.6.1.dev9
+Version: 0.6.1.dev13
 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.6.1.dev9 → jaxsim-0.6.1.dev13}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.6.1.dev9'
-__version_tuple__ = version_tuple = (0, 6, 1, 'dev9')
+__version__ = version = '0.6.1.dev13'
+__version_tuple__ = version_tuple = (0, 6, 1, 'dev13')

{jaxsim-0.6.1.dev9 → jaxsim-0.6.1.dev13}/src/jaxsim/math/rotation.py

@@ -68,26 +68,17 @@ class Rotation:
 
         vector = vector.squeeze()
 
-        def theta_is_not_zero(axis: jtp.Vector) -> jtp.Matrix:
+        theta = safe_norm(vector)
 
-            v = axis
-            theta = safe_norm(v)
+        s = jnp.sin(theta)
+        c = jnp.cos(theta)
 
-            s = jnp.sin(theta)
-            c = jnp.cos(theta)
+        c1 = 2 * jnp.sin(theta / 2.0) ** 2
 
-            c1 = 2 * jnp.sin(theta / 2.0) ** 2
+        safe_theta = jnp.where(theta == 0, 1.0, theta)
+        u = vector / safe_theta
+        u = jnp.vstack(u.squeeze())
 
-            u = v / theta
-            u = jnp.vstack(u.squeeze())
+        R = c * jnp.eye(3) - s * Skew.wedge(u) + c1 * u @ u.T
 
-            R = c * jnp.eye(3) - s * Skew.wedge(u) + c1 * u @ u.T
-
-            return R.transpose()
-
-        return jnp.where(
-            jnp.allclose(vector, 0.0),
-            # Return an identity rotation matrix when the input vector is zero.
-            jnp.eye(3),
-            theta_is_not_zero(axis=vector),
-        )
+        return R.transpose()

{jaxsim-0.6.1.dev9 → jaxsim-0.6.1.dev13}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: jaxsim
-Version: 0.6.1.dev9
+Version: 0.6.1.dev13
 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.6.1.dev9 → jaxsim-0.6.1.dev13}/tests/test_benchmark.py

@@ -9,13 +9,14 @@ import jaxsim.api as js
 
 def vectorize_data(model: js.model.JaxSimModel, batch_size: int):
     key = jax.random.PRNGKey(seed=0)
+    keys = jax.random.split(key, num=batch_size)
 
     return jax.vmap(
         lambda key: js.data.random_model_data(
             model=model,
             key=key,
         )
-    )(jax.numpy.repeat(key[None, :], repeats=batch_size, axis=0))
+    )(keys)
 
 
 def benchmark_test_function(
@@ -26,7 +27,10 @@ def benchmark_test_function(
 
     # Warm-up call to avoid including compilation time
     jax.vmap(func, in_axes=(None, 0))(model, data)
-    benchmark(jax.vmap(func, in_axes=(None, 0)), model, data)
+
+    # Benchmark the function call
+    # Note: jax.block_until_ready is used to ensure that the benchmark is not measuring only the asynchronous dispatch
+    benchmark(jax.block_until_ready(jax.vmap(func, in_axes=(None, 0))), model, data)
 
 
 @pytest.mark.benchmark