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)

        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.
        
dpkd/transformers/src/transformers/models/rag/modeling_tf_rag.py