docstr: This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
  [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
  beam_idx at every generation step.
docstr_tok: ["This","function","is","used","to","re-order","the","`","past_key_values","`","cache","if","[","`","~PretrainedModel.beam_search","`","]","or","[","`","~PretrainedModel.beam_sample","`","]","is","called",".","This","is","required","to","match","`","past_key_values","`","with","the","correct","beam_idx","at","every","generation","step","."]
code: def _reorder_cache(
      self,
      past_key_values: Tuple[Tuple[torch.Tensor]],
      beam_idx: torch.Tensor,
  ) -> Tuple[Tuple[torch.Tensor]]:
      """
      This function is used to re-order the `past_key_values` cache if [`~PretrainedModel.beam_search`] or
      [`~PretrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
      beam_idx at every generation step.
      """
      return self._default._reorder_cache(past_key_values, beam_idx)
  
imports: import argparse
  import inspect
  import logging
  from typing import Tuple
  import torch
  from accelerate import PartialState
  from accelerate.utils import set_seed
  from transformers import AutoTokenizer, BloomForCausalLM, BloomTokenizerFast, CTRLLMHeadModel, CTRLTokenizer, GenerationMixin, GPT2LMHeadModel, GPT2Tokenizer, GPTJForCausalLM, LlamaForCausalLM, LlamaTokenizer, OpenAIGPTLMHeadModel, OpenAIGPTTokenizer, OPTForCausalLM, TransfoXLLMHeadModel, TransfoXLTokenizer, XLMTokenizer, XLMWithLMHeadModel, XLNetLMHeadModel, XLNetTokenizer
  from transformers.modeling_outputs import CausalLMOutputWithPast
label_desc: Use image node_id 5 for calling the _ModelFallbackWrapper obj's underlying member method code with example usage: obj._reorder_cache(past_key_values, beam_idx) and returns: self