code: def _gather(  # pylint: disable=unused-argument
      params,
      indices,
      validate_indices=None,
      axis=None,
      batch_dims=0,
      name=None,
  ):
      """gather."""
      params = ops.convert_to_tensor(params)
      indices = ops.convert_to_tensor(indices, dtype_hint=np.int32)
      if validate_indices is not None:
          raise NotImplementedError(
              "Argument `validate_indices != None` is currently unimplemented."
          )
      if batch_dims < 0:
          raise NotImplementedError(
              "Negative `batch_dims` is currently unsupported."
          )
      if axis is None:
          axis = batch_dims
      if axis < 0:
          axis = axis + len(params.shape)
      # NOTE: For only the numpy backend, this function could create a single result
      # ndarray and use in-place updates.  For the Jax backend, this function
      # vmaps `np.take`.
      if JAX_MODE:
          if batch_dims == 0 and axis == 0:
              return params[indices]
          take = lambda params, indices: np.take(
              params,
              indices,  # pylint: disable=g-long-lambda
              axis=axis - batch_dims,
          )
          take = functools.reduce(
              lambda g, f: f(g), [jax.vmap] * int(batch_dims), take
          )
          return take(params, indices)
      params = ops.convert_to_tensor(params)
      res = np.array(
          [
              np.take(params[i], indices[i], axis=axis - batch_dims)
              for i in np.ndindex(*params.shape[:batch_dims])
          ]
      )
      return np.reshape(
          res,
          params.shape[:axis]
          + indices.shape[batch_dims:]
          + params.shape[axis + 1 :],
      )
  
imports: import functools
  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.linalg_impl import einsum
  from tensorflow_probability.python.internal.backend.numpy.linalg_impl import norm
  from tensorflow_probability.python.internal.backend.numpy.linalg_impl import tensordot
label_desc: Use image node_id 2 for calling a global function with example usage: _gather(params, indices, validate_indices, axis, batch_dims, name) and returns: np, take, params