code: def _ensure_jax_install():  # pylint: disable=g-statement-before-imports
      """Attempt to import JAX.
  
      Raises:
        ImportError: if JAX is not importable.
      """
      try:
          import jax
  
          del jax
      except ImportError:
          print(
              "\n\nFailed to import JAX. " "pip install jax jaxlib\n\n"
          )
          raise
  
file_path: probability/tensorflow_probability/substrates/jax/__init__.py
imports: from tensorflow_probability.python.version import __version__
  from tensorflow_probability.substrates.jax import bijectors
  from tensorflow_probability.substrates.jax import distributions
  from tensorflow_probability.substrates.jax import experimental
  from tensorflow_probability.substrates.jax import glm
  from tensorflow_probability.substrates.jax import internal
  from tensorflow_probability.substrates.jax import math
  from tensorflow_probability.substrates.jax import mcmc
  from tensorflow_probability.substrates.jax import monte_carlo
  from tensorflow_probability.substrates.jax import optimizer
  from tensorflow_probability.substrates.jax import random
  from tensorflow_probability.substrates.jax import stats
  from tensorflow_probability.substrates.jax import sts
  from tensorflow_probability.substrates.jax import util
  from tensorflow_probability.substrates.jax import vi
  from tensorflow_probability.python.internal.backend import jax
label_desc: Use image node_id 1 for calling a global function with example usage: _ensure_jax_install() without return types