code: def _expand_inputs_for_generation(
      input_ids: torch.Tensor,
      attention_mask: torch.Tensor,
      last_hidden_state: torch.Tensor,
      expand_size: int = 1,
  ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
      expanded_return_idx = (
          torch.arange(input_ids.shape[0])
          .view(-1, 1)
          .repeat(1, expand_size)
          .view(-1)
          .to(input_ids.device)
      )
      input_ids = input_ids.index_select(0, expanded_return_idx)
  
      attention_mask = attention_mask.index_select(
          0, expanded_return_idx
      )
  
      last_hidden_state = last_hidden_state.index_select(
          0, expanded_return_idx.to(last_hidden_state.device)
      )
      return input_ids, attention_mask, last_hidden_state
  
imports: import copy
  import itertools
  from typing import List, Optional, Tuple
  import torch
  import torch.nn.functional
  from transformers import BartConfig
  from transformers.generation_utils import GenerationMixin
label_desc: Use image node_id 2 for calling the BARTBeamSearchGenerator obj's underlying member method code with example usage: obj._expand_inputs_for_generation(input_ids, attention_mask, last_hidden_state, expand_size) and returns: input_ids,  attention_mask,  last_hidden_state