src/transformers/models/roformer/modeling_roformer.py/RoFormerForCausalLM/_reorder_cache
class RoFormerForCausalLM:    def _reorder_cache(self, past, beam_idx):
        reordered_past = ()
        for layer_past in past:
            reordered_past += (
                tuple(past_state.index_select(0, beam_idx) for past_state in layer_past[:2]) + layer_past[2:],
            )
        return reordered_past