src/transformers/models/rag/modeling_tf_rag.py/TFRagTokenForGeneration/_reorder_cache
class TFRagTokenForGeneration:    def _reorder_cache(past, beam_idx):
        """Reorders cache for generation. BART-inspired but we need to take care of the extra dimension for docs"""

        def tf_index_select(input_, dim, indices):
            """
            Input:
                input_(tensor): input tensor dim(int): dimension indices(list): selected indices list
            Output:
                mimic of torch_tensor.index_select(dim, indices)

            credit: https://stackoverflow.com/questions/58464790/is-there-an-equivalent-function-of-pytorch-named-index-select-in-tensorflow
            """
            shape = shape_list(input_)
            if dim == -1:
                dim = len(shape) - 1
            shape[dim] = 1

            tmp = []
            for idx in indices:
                begin = [0] * len(shape)
                begin[dim] = idx
                tmp.append(tf.slice(input_, begin, shape))
            res = tf.concat(tmp, axis=dim)

            return res

        def _reorder_stacked(hidden_states, new_order=beam_idx):
            n_docs = hidden_states.shape[0] // new_order.shape[0]
            hidden_states = tf.reshape(hidden_states, (-1, n_docs, *hidden_states.shape[1:]))
            hidden_states = tf_index_select(hidden_states, 0, new_order)
            return tf.reshape(hidden_states, (-1, *hidden_states.shape[2:]))

        if len(past) == 1:
            return past

        past_key_values = past[1]

        reordered_past = ()
        for layer_past in past_key_values:
            reordered_past += (tuple(_reorder_stacked(past_state, beam_idx) for past_state in layer_past),)

        return (past[0], reordered_past)