Make JAX backend import lazy.

We observed the following surprising behavior: importing tensorflow also imports jax, if installed. This happens because tensorflow has an opt_einsum dependency that is imported eagerly, and the opt_einsum jax backend eagerly imports jax. To avoid this, make the jax import from opt_einsum lazy.

opt_einsum/backends/jax.py

@@ -8,22 +8,29 @@
 
 __all__ = ["build_expression", "evaluate_constants"]
 
-try:
+
+_JAX = None
+
+
+def _get_jax_and_to_jax():
+  global _JAX
+  if _JAX is None:
     import jax
 
     @to_backend_cache_wrap
     @jax.jit
     def to_jax(x):
         return x
 
-except ImportError:
-    pass
+    _JAX = jax, to_jax
+
+  return _JAX
 
 
 def build_expression(_, expr):  # pragma: no cover
     """Build a jax function based on ``arrays`` and ``expr``.
     """
-    import jax
+    jax, _ = _get_jax_and_to_jax()
 
     jax_expr = jax.jit(expr._contract)
 
@@ -37,4 +44,6 @@ def evaluate_constants(const_arrays, expr):  # pragma: no cover
     """Convert constant arguments to jax arrays, and perform any possible
     constant contractions.
     """
+    jax, to_jax = _get_jax_and_to_jax()
+
     return expr(*[to_jax(x) for x in const_arrays], backend='jax', evaluate_constants=True)