jaxsim 0.3.1.dev17__py3-none-any.whl → 0.3.1.dev21__py3-none-any.whl

jaxsim/_version.py

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

jaxsim/exceptions.py

@@ -0,0 +1,63 @@
+import jax
+
+
+def raise_if(
+    condition: bool | jax.Array, exception: type, msg: str, *args, **kwargs
+) -> None:
+    """
+    Raise a host-side exception if a condition is met. Useful in jit-compiled functions.
+
+    Args:
+        condition:
+            The boolean condition of the evaluated expression that triggers
+            the exception during runtime.
+        exception: The type of exception to raise.
+        msg:
+            The message to display when the exception is raised. The message can be a
+            format string (fmt), whose fields are filled with the args and kwargs.
+    """
+
+    # Check early that the format string is well-formed.
+    try:
+        _ = msg.format(*args, **kwargs)
+    except Exception as e:
+        msg = "Error in formatting exception message with args={} and kwargs={}"
+        raise ValueError(msg.format(args, kwargs)) from e
+
+    def _raise_exception(condition: bool, *args, **kwargs) -> None:
+        """The function called by the JAX callback."""
+
+        if condition:
+            raise exception(msg.format(*args, **kwargs))
+
+    def _callback(args, kwargs) -> None:
+        """The function that calls the JAX callback, executed only when needed."""
+
+        jax.debug.callback(_raise_exception, condition, *args, **kwargs)
+
+    # Since running a callable on the host is expensive, we prevent its execution
+    # if the condition is False with a low-level conditional expression.
+    def _run_callback_only_if_condition_is_true(*args, **kwargs) -> None:
+        return jax.lax.cond(
+            condition,
+            _callback,
+            lambda args, kwargs: None,
+            args,
+            kwargs,
+        )
+
+    return _run_callback_only_if_condition_is_true(*args, **kwargs)
+
+
+def raise_runtime_error_if(
+    condition: bool | jax.Array, msg: str, *args, **kwargs
+) -> None:
+
+    return raise_if(condition, RuntimeError, msg, *args, **kwargs)
+
+
+def raise_value_error_if(
+    condition: bool | jax.Array, msg: str, *args, **kwargs
+) -> None:
+
+    return raise_if(condition, ValueError, msg, *args, **kwargs)

{jaxsim-0.3.1.dev17.dist-info → jaxsim-0.3.1.dev21.dist-info}/METADATA

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

{jaxsim-0.3.1.dev17.dist-info → jaxsim-0.3.1.dev21.dist-info}/RECORD

@@ -1,5 +1,6 @@
 jaxsim/__init__.py,sha256=xzuTuZrgKdWLqqDzbvqzm2cJrEtAbepOeUqDu7ByVek,2621
-jaxsim/_version.py,sha256=EQQfkY5WXMHFjdRnYAQqABGWC0VK4dlpuNh_wr1KxYA,426
+jaxsim/_version.py,sha256=fSekabX0ZEIHwkdp0Sa0iQ2H7hhPzvja3dhu7EFiX4I,426
+jaxsim/exceptions.py,sha256=8_h8iqL8DgNR754dR8SZiQ7361GR5V1sUk3ZuZCHw1Q,2069
 jaxsim/logging.py,sha256=c4zhwBKf9eAYAHVp62kTEllqdsZgh0K-kPKVy8L3elU,1584
 jaxsim/typing.py,sha256=cl7HHQCeP3mHmtF6EuQZcCjGvDmc_AryMWntP_lRBGg,722
 jaxsim/api/__init__.py,sha256=8eV22t2S3UwNyCg8karPetG1dmX1VDBXkyv28_FwNQA,210
@@ -58,8 +59,8 @@ jaxsim/utils/__init__.py,sha256=Y5zyoRevl3EMVQadhZ4EtSwTEkDt2vcnFoRhPJjKTZ0,215
 jaxsim/utils/jaxsim_dataclass.py,sha256=h26timZ_XrBL_Q_oymv-DkQd-EcUiHn8QexAaZXBY9c,11396
 jaxsim/utils/tracing.py,sha256=KDMoyVPlu2NJvFkhtZwq5AkqMMgajt3munvJom-vEjQ,650
 jaxsim/utils/wrappers.py,sha256=QIJitSoljrKR_U4T3ewCJPT3DTh-tPZsRsg0t_MH93E,3896
-jaxsim-0.3.1.dev17.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
-jaxsim-0.3.1.dev17.dist-info/METADATA,sha256=zRsMl96hDJt919NgrEuxkhye1S8X20bi_nWdPzJiptU,9739
-jaxsim-0.3.1.dev17.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-jaxsim-0.3.1.dev17.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
-jaxsim-0.3.1.dev17.dist-info/RECORD,,
+jaxsim-0.3.1.dev21.dist-info/LICENSE,sha256=eaYdFmdeMbiIoIiPzEK0MjP1S9wtFXjXNR5er49uLR0,1546
+jaxsim-0.3.1.dev21.dist-info/METADATA,sha256=wtxQdWa5FFEqYdZx81i-VgNk7DKBY6YMQAXn5_1ctMY,9739
+jaxsim-0.3.1.dev21.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+jaxsim-0.3.1.dev21.dist-info/top_level.txt,sha256=LxGMA8FLtXjQ6oI7N5gd_R_oSUHxpXxUEOfT1xS_ni0,7
+jaxsim-0.3.1.dev21.dist-info/RECORD,,