Fix check import jax

fluiddyn · Aug 17, 2024 · 1b7f178 · 1b7f178
1 parent b9ee14b
commit 1b7f178
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transonic/util.py b/src/transonic/util.py
@@ -131,7 +131,7 @@ def can_import_accelerator(backend: str = backend_default):
     elif backend == "jax":
         try:
             import jax
-        except ImportError:
+        except (ImportError, AttributeError, RuntimeError):
             return False
     elif backend == "python":
         return True