code: def _gamma_jax(
      shape, alpha, beta=None, dtype=np.float32, seed=None, name=None
  ):  # pylint: disable=unused-argument
      """JAX-based reparameterized gamma sampler."""
      dtype = utils.common_dtype([alpha, beta], dtype_hint=dtype)
      alpha = np.array(alpha, dtype=dtype)
      beta = None if beta is None else np.array(beta, dtype=dtype)
      shape = _ensure_shape_tuple(shape)
      import jax.random as jaxrand  # pylint: disable=g-import-not-at-top
  
      if seed is None:
          raise ValueError("Must provide PRNGKey to sample in JAX.")
      # TODO(srvasude): Sample in the given dtype once
      # https://github.com/google/jax/issues/2130 is fixed.
      samps = jaxrand.gamma(
          key=seed, a=alpha, shape=shape, dtype=np.float64
      ).astype(dtype)
      # Match the 0->tiny behavior of tf.random.gamma.
      return np.maximum(
          np.finfo(dtype).tiny, samps if beta is None else samps / beta
      )
  
imports: import numpy
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import ops
  from tensorflow_probability.python.internal.backend.numpy.numpy_math import softmax
label_desc: Use image node_id 7 for calling a global function with example usage: _gamma_jax(shape, alpha, beta, dtype, seed, name) and returns: np