jaxsim 0.3.1.dev46__tar.gz → 0.3.1.dev51__tar.gz

{jaxsim-0.3.1.dev46 → jaxsim-0.3.1.dev51}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.3.1.dev46
+Version: 0.3.1.dev51
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it

{jaxsim-0.3.1.dev46 → jaxsim-0.3.1.dev51}/src/jaxsim/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.3.1.dev46'
-__version_tuple__ = version_tuple = (0, 3, 1, 'dev46')
+__version__ = version = '0.3.1.dev51'
+__version_tuple__ = version_tuple = (0, 3, 1, 'dev51')

{jaxsim-0.3.1.dev46 → jaxsim-0.3.1.dev51}/src/jaxsim/api/model.py

@@ -1841,7 +1841,13 @@ def step(
         and the new state of the integrator.
     """
 
-    integrator_kwargs = kwargs if kwargs is not None else dict()
+    # Extract the integrator kwargs.
+    # The following logic allows using integrators having kwargs colliding with the
+    # kwargs of this step function.
+    kwargs = kwargs if kwargs is not None else {}
+    integrator_kwargs = kwargs.pop("integrator_kwargs", {})
+    integrator_kwargs = kwargs | integrator_kwargs
+
     integrator_state = integrator_state if integrator_state is not None else dict()
 
     # Extract the initial resources.
@@ -1855,8 +1861,21 @@ def step(
         t0=jnp.array(t0_ns / 1e9).astype(float),
         dt=dt,
         params=integrator_state_x0,
+        # Always inject the current (model, data) pair into the system dynamics
+        # considered by the integrator, and include the input variables represented
+        # by the pair (joint_forces, link_forces).
+        # Note that the wrapper of the system dynamics will override (state_x0, t0)
+        # inside the passed data even if it is not strictly needed. This logic is
+        # necessary to re-use the jit-compiled step function of compatible pytrees
+        # of model and data produced e.g. by parameterized applications.
         **(
-            dict(joint_forces=joint_forces, link_forces=link_forces) | integrator_kwargs
+            dict(
+                model=model,
+                data=data,
+                joint_forces=joint_forces,
+                link_forces=link_forces,
+            )
+            | integrator_kwargs
         ),
     )
 

{jaxsim-0.3.1.dev46 → jaxsim-0.3.1.dev51}/src/jaxsim/integrators/common.py

@@ -422,7 +422,9 @@ class ExplicitRungeKutta(Integrator[PyTreeType, PyTreeType], Generic[PyTreeType]
 
         # Update the FSAL property for the next iteration.
         if self.has_fsal:
-            self.params["dxdt0"] = jax.tree_map(lambda l: l[self.index_of_fsal], K)
+            self.params["dxdt0"] = jax.tree_util.tree_map(
+                lambda l: l[self.index_of_fsal], K
+            )
 
         # Compute the output state.
         # Note that z contains as many new states as the rows of `b.T`.

{jaxsim-0.3.1.dev46 → jaxsim-0.3.1.dev51}/src/jaxsim.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxsim
-Version: 0.3.1.dev46
+Version: 0.3.1.dev51
 Home-page: https://github.com/ami-iit/jaxsim
 Author: Diego Ferigo
 Author-email: diego.ferigo@iit.it

{jaxsim-0.3.1.dev46 → jaxsim-0.3.1.dev51}/tests/test_exceptions.py

@@ -68,21 +68,21 @@ def test_exceptions_in_jit_functions():
     assert jit_compiled_function._cache_size() == 1
 
     # Let's trigger a ValueError exception by passing 42.
-    # Note: the real ValueError is printed in a stream that I couldn't figure out
-    # how to capture in pytest.
-    with pytest.raises(jaxlib.xla_extension.XlaRuntimeError):
-
-        data = 42
+    data = 42
+    with pytest.raises(
+        jaxlib.xla_extension.XlaRuntimeError,
+        match=f"ValueError: Raising ValueError since data={data}",
+    ):
         _ = jit_compiled_function(data=data)
 
     assert jit_compiled_function._cache_size() == 1
 
     # Let's trigger a RuntimeError exception by passing -42.
-    # Note: the real RuntimeError is printed in a stream that I couldn't figure out
-    # how to capture in pytest.
-    with pytest.raises(jaxlib.xla_extension.XlaRuntimeError):
-
-        data = -42
+    data = -42
+    with pytest.raises(
+        jaxlib.xla_extension.XlaRuntimeError,
+        match=f"RuntimeError: Raising RuntimeError since data={data}",
+    ):
         _ = jit_compiled_function(data=data)
 
     assert jit_compiled_function._cache_size() == 1