Disable host callbacks if no MLIR translation rule is defined

ami-iit · Jan 5, 2025 · 356a33d · 356a33d
1 parent 8382d68
commit 356a33d
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/jaxsim/exceptions.py b/src/jaxsim/exceptions.py
@@ -21,9 +21,10 @@ def raise_if(
         **kwargs: The keyword arguments to fill the format string
     """
 
-    # Disable host callback if running on TPU.
-    if jax.devices()[0].platform == "tpu" or os.environ.get(
-        "JAXSIM_DISABLE_EXCEPTIONS", 0
+    # Disable host callback if running on unsupported hardware or if the user
+    # explicitly disabled it.
+    if jax.devices()[0].platform in {"tpu", "METAL"} or os.environ.get(
+        "JAXSIM_DISABLE_EXCEPTIONS", "0"
     ):
         return