docstr: Used by HuggingFace generate when using beam search with kv-caching.
  
  See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133
  for an example in transformers.
docstr_tok: ["Used","by","HuggingFace","generate","when","using","beam","search","with","kv-caching",".","See","https",":","\/\/github.com\/huggingface\/transformers\/blob\/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec\/src\/transformers\/models\/gpt2\/modeling_gpt2.py","#","L1122-L1133","for","an","example","in","transformers","."]
code: def _reorder_cache(past_key_values, beam_idx):
      """Used by HuggingFace generate when using beam search with kv-caching.
  
      See https://github.com/huggingface/transformers/blob/3ec7a47664ebe40c40f4b722f6bb1cd30c3821ec/src/transformers/models/gpt2/modeling_gpt2.py#L1122-L1133
      for an example in transformers.
      """
      reordered_past = []
      for layer_past in past_key_values:
          reordered_past += [
              tuple(
                  (
                      past_state.index_select(0, beam_idx)
                      for past_state in layer_past
                  )
              )
          ]
      return reordered_past
  
imports: import math
  import os
  import warnings
  from typing import List, Optional, Tuple, Union
  import torch
  import torch.nn
  import torch.nn.functional
  from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
  from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
  from einops import rearrange
  from packaging import version
  from text_generation_server.utils.layers import TensorParallelEmbedding, TensorParallelColumnLinear, TensorParallelRowLinear, TensorParallelHead, get_linear
label_desc: Use image node_id 4 for calling the MPTForCausalLM obj's underlying member method code with example usage: obj._reorder_cache(past_key_values, beam_idx) and returns: reordered_past