Skip to content

Commit

Permalink
Soften need for jax, np, or tf.
Browse files Browse the repository at this point in the history
Many folks use edward2 just for the robust_map utility.

PiperOrigin-RevId: 630494865
  • Loading branch information
dustinvtran authored and edward-bot committed May 3, 2024
1 parent 83607d9 commit 9db6c50
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions edward2/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,14 +21,14 @@
https://github.com/google/edward2/blob/main/README.md)
+ [Upgrading from Edward to Edward2](
https://github.com/google/edward2/blob/main/Upgrading_From_Edward_To_Edward2.md)
"""

import warnings

__all__ = []
try:
from edward2 import jax

__all__ += ["jax"]
except (ImportError, AttributeError):
# We also catch AttributeError. JAX can raise AttributeError during
Expand All @@ -37,6 +37,7 @@

try:
from edward2 import numpy

__all__ += ["numpy"]
except ImportError:
warnings.warn("NumPy backend for Edward2 is not available.")
Expand All @@ -51,6 +52,8 @@


if all(x not in __all__ for x in ("jax", "numpy", "tensorflow")):
raise ImportError("No backend is available for Edward2. Please install "
"dependencies for JAX, NumPy, or TensorFlow.")
warnings.warn(
"No backend is available for Edward2. Please install dependencies for"
" JAX, NumPy, or TensorFlow."
)
# pylint: enable=g-import-not-at-top

0 comments on commit 9db6c50

Please sign in to comment.