Catch AttributeError during JAX import.

Importing edward2 leads to AttributeError if you use a Colab runtime without JAX dependency. This fix catches AttributeError and omits the edward2.jax namespace. PiperOrigin-RevId: 548798730
google · Oct 16, 2023 · 636704d · 636704d
1 parent 3a9b7ef
commit 636704d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/edward2/__init__.py b/edward2/__init__.py
@@ -30,7 +30,9 @@
 try:
   from edward2 import jax
   __all__ += ["jax"]
-except ImportError:
+except (ImportError, AttributeError):
+  # We also catch AttributeError. JAX can raise AttributeError during
+  # imports in colab unless JAX is directly installed first.
   warnings.warn("JAX backend for Edward2 is not available.")
 
 try: