code: def _gather_nd(  # pylint: disable=unused-argument
      params, indices, batch_dims=0, name=None
  ):
      """gather_nd."""
      params = ops.convert_to_tensor(params)
      indices = ops.convert_to_tensor(indices, dtype_hint=np.int32)
      if batch_dims < 0:
          raise NotImplementedError(
              "Negative `batch_dims` is currently unsupported."
          )
      if not JAX_MODE and batch_dims > 0:
          raise NotImplementedError(
              "`batch_dims > 0` currently unsupported in NumPy backend."
          )
      gather_nd_ = _gather_nd_single
      if JAX_MODE:
          gather_nd_ = functools.reduce(
              lambda g, f: f(g),
              [jax.vmap] * int(batch_dims),
              gather_nd_,
          )
      return gather_nd_(params, indices)
  
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 6 for calling a global function with example usage: _gather_nd(params, indices, batch_dims, name) and returns: gather_nd_