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.
  
  Output shares the same memory storage as `past`.
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",".","Output","shares","the","same","memory","storage","as","`","past","`","."]
code: def _reorder_cache(
      self,
      past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...],
      beam_idx: torch.LongTensor,
  ) -> Tuple[Tuple[torch.Tensor, 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.
  
      Output shares the same memory storage as `past`.
      """
      # Get a copy of `beam_idx` on all the devices where we need those indices.
      device_to_beam_idx = {
          past_state.device: beam_idx.to(past_state.device)
          for layer_past in past
          for past_state in layer_past
      }
      reordered_past = tuple(
          (
              layer_past[0].index_select(
                  0, device_to_beam_idx[layer_past[0].device]
              ),
              layer_past[1].index_select(
                  0, device_to_beam_idx[layer_past[0].device]
              ),
          )
          for layer_past in past
      )
      return reordered_past
  
imports: import math
  from typing import Optional, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
  from torch.nn import functional
  from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
  from ...modeling_attn_mask_utils import _prepare_4d_causal_attention_mask
  from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions, CausalLMOutputWithCrossAttentions, QuestionAnsweringModelOutput, SequenceClassifierOutputWithPast, TokenClassifierOutput
  from ...modeling_utils import PreTrainedModel
  from ...utils import logging
  from .configuration_mpt import MptConfig
label_desc: Use image node_id 6 for calling the MptForCausalLM obj's underlying member method code with example usage: obj._reorder_cache(past, beam_idx) and returns: reordered_past