From 2e53f6c98cb6d544ae5a24a2aa661fc18d2a886a Mon Sep 17 00:00:00 2001 From: Dustin Tran Date: Mon, 17 Jul 2023 14:22:44 -0700 Subject: [PATCH] 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 --- edward2/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/edward2/__init__.py b/edward2/__init__.py index 0a865b3d..4e3b05b1 100644 --- 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: