code: def _random_gamma_grad(
      alpha, sample, name=None
  ):  # pylint: disable=unused-argument
      if JAX_MODE:
          import jax.lax  # pylint: disable=g-import-not-at-top
  
          # Broadcast alpha and samples.
          broadcast_shape = (alpha + sample).shape
          alpha = np.broadcast_to(alpha, broadcast_shape)
          sample = np.broadcast_to(sample, broadcast_shape)
          return jax.lax.random_gamma_grad(alpha, sample)
      raise NotImplementedError
  
imports: import collections
  import numpy
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
label_desc: Use image node_id 3 for calling a global function with example usage: _random_gamma_grad(alpha, sample, name) and returns: jax