code: def _axis_gather(params, indices, axis):
      """Helper that implements ragged gather when axis>0 and batch_dims==0.
  
      Args:
        params: The tensor from which to gather values.
        indices: The indices of values to gather.
        axis: The axis in `params` to gather `indices` from.
  
      Returns:
        A potentially ragged tensor.
      """
      if axis > 1:
          if not isinstance(params, ragged_tensor.RaggedTensor):
              params = ragged_tensor.RaggedTensor.from_tensor(
                  params,
                  ragged_rank=1,
                  row_splits_dtype=indices.row_splits.dtype,
              )
          # Recurse, using the flattened params (but do not flatten indices).
          return params.with_values(
              _gather(params.values, indices, axis - 1, 0)
          )
  
      if indices.shape.rank is None:
          raise ValueError("rank(indices) must be known statically")
  
      # Note: there is no checking of indices. If there is some index
      # out of bounds, the results may be nonsensical.
  
      assert axis == 1
      # If params.shape=[P1...PN] and indices.shape=[I1...IM], then:
      #
      #     output[p1,      i1...im,     p3...pn] =
      #     params[p1, indices[i1...im], p3...pn]
      #
      # We construct `output` by flattening `params`, adjusting the `indices` to
      # have one additional dimension, and to point into that flattened list, and
      # recursively calling `gather`.
      flat_params = _flatten_dims_0_and_1(params)
      adjustments = _row_starts(
          params, indices.dtype
      )  # offset for each batch
      adjustments = _increase_rank_to(
          adjustments, indices.shape.ndims + 1
      )
      adjusted_indices = indices + adjustments
      return _gather(flat_params, adjusted_indices, axis - 1, 0)
  
imports: from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import indexed_slices
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import gen_ragged_array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops.ragged import ragged_array_ops
  from tensorflow.python.ops.ragged import ragged_math_ops
  from tensorflow.python.ops.ragged import ragged_tensor
  from tensorflow.python.util import dispatch
label_desc: Use image node_id 4 for calling a global function with example usage: _axis_gather(params, indices, axis) and returns: _gather, params