docstr: RAG-specific `_gather_beams`: gathers the beam slices indexed by beam_indices into new beam array. If the
  nested tensor has a shape mismatch with the beam indices, then it means it is the cache. In that case, isolates
  and takes care of the extra dimension for ndocs.
docstr_tok: ["RAG-specific","`","_gather_beams","`",":","gathers","the","beam","slices","indexed","by","beam_indices","into","new","beam","array",".","If","the","nested","tensor","has","a","shape","mismatch","with","the","beam","indices",",","then","it","means","it","is","the","cache",".","In","that","case",",","isolates","and","takes","care","of","the","extra","dimension","for","ndocs","."]
code: def _gather_beams(nested, beam_indices, batch_axis=0):
      """
      RAG-specific `_gather_beams`: gathers the beam slices indexed by beam_indices into new beam array. If the
      nested tensor has a shape mismatch with the beam indices, then it means it is the cache. In that case, isolates
      and takes care of the extra dimension for ndocs.
      """
  
      def gather_fn(tensor):
          is_rag_cache = tensor.shape[0] != beam_indices.shape[0]
          if is_rag_cache:
              n_docs = tensor.shape[0] // beam_indices.shape[0]
              batch_size = beam_indices.shape[0]
              # reshapes into (batch size, num beams, n_docs, ...), the cache format expected by RAG
              tensor = tf.reshape(
                  tensor, (batch_size, -1, n_docs, *tensor.shape[2:])
              )
  
          gathered_tensor = tf.gather(
              params=tensor, indices=beam_indices, axis=1, batch_dims=1
          )
  
          if is_rag_cache:
              # reshapes back into the shape expected by beam search
              gathered_tensor = tf.reshape(
                  gathered_tensor,
                  (batch_size * n_docs, -1, *gathered_tensor.shape[3:]),
              )
  
          return gathered_tensor
  
      return tf.nest.map_structure(gather_fn, nested)
  
imports: from __future__ import annotations
  import copy
  from dataclasses import dataclass
  from typing import List, Optional, Tuple, Union
  import numpy
  import tensorflow
  from ...configuration_utils import PretrainedConfig
  from ...generation import TFLogitsProcessorList
  from ...modeling_tf_utils import TFCausalLanguageModelingLoss, TFModelInputType, TFPreTrainedModel, shape_list, unpack_inputs
  from ...utils import ModelOutput, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from .configuration_rag import RagConfig
  from .retrieval_rag import RagRetriever
label_desc: Use image node_id 7 for calling the TFRagTokenForGeneration obj's underlying member method code with example usage: obj._gather_beams(nested, beam_indices, batch_axis) and returns: tf, gathered_tensor