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 · Jul 17, 2023 · 2e53f6c · 2e53f6c
1 parent 7a4e1df
commit 2e53f6c
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: