code: def ViTModel_pipeline_forward(
      stage_manager: PipelineStageManager, stage_index: List[int]
  ):
      from transformers.models.vit.modeling_vit import (
          BaseModelOutputWithPooling,
      )
  
      def pp_forward(
          self,
          pixel_values: Optional[torch.Tensor] = None,
          bool_masked_pos: Optional[torch.BoolTensor] = None,
          head_mask: Optional[torch.Tensor] = None,
          output_attentions: Optional[bool] = None,
          output_hidden_states: Optional[bool] = None,
          interpolate_pos_encoding: Optional[bool] = None,
          return_dict: Optional[bool] = None,
          hidden_states: Optional[torch.FloatTensor] = None,
      ) -> Union[Tuple, BaseModelOutputWithPooling]:
          r"""
          bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*):
              Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
          """
          return_dict = (
              return_dict
              if return_dict is not None
              else self.config.use_return_dict
          )
  
          logger = logging.get_logger(__name__)
  
          # Preprocess passed in arguments
          if output_attentions:
              logger.warning_once(
                  "output_attentions=True is not supported for pipeline models at the moment."
              )
              output_attentions = False
          if output_hidden_states:
              logger.warning_once(
                  "output_hidden_states=True is not supported for pipeline models at the moment."
              )
              output_hidden_states = False
  
          # Prepare head mask if needed
          # 1.0 in head_mask indicate we keep the head
          # attention_probs has shape bsz x n_heads x N x N
          # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
          # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
          head_mask = self.get_head_mask(
              head_mask, self.config.num_hidden_layers
          )
  
          if stage_manager.is_first_stage():
              if pixel_values is None:
                  raise ValueError("You have to specify pixel_values")
  
              # TODO(FoolPlayer): maybe have a cleaner way to cast the input (from `ImageProcessor` side?)
              expected_dtype = (
                  self.embeddings.patch_embeddings.projection.weight.dtype
              )
              if pixel_values.dtype != expected_dtype:
                  pixel_values = pixel_values.to(expected_dtype)
  
              embedding_output = self.embeddings(
                  pixel_values,
                  bool_masked_pos=bool_masked_pos,
                  interpolate_pos_encoding=interpolate_pos_encoding,
              )
              hidden_states = embedding_output
          else:
              assert (
                  hidden_states is not None
              ), f"Current stage is {stage_manager.stage}, hidden_states should not be None"
  
          encoder_outputs = _encoder_forward(
              encoder=self.encoder,
              start_idx=stage_index[0],
              end_idx=stage_index[1],
              hidden_states=hidden_states,
              head_mask=head_mask,
              return_dict=return_dict,
              stage_manager=stage_manager,
          )
          if not stage_manager.is_last_stage():
              return {"hidden_states": encoder_outputs}
  
          sequence_output = encoder_outputs[0]
          sequence_output = self.layernorm(sequence_output)
          pooled_output = (
              self.pooler(sequence_output)
              if self.pooler is not None
              else None
          )
  
          if not return_dict:
              head_outputs = (
                  (sequence_output, pooled_output)
                  if pooled_output is not None
                  else (sequence_output,)
              )
              return head_outputs + encoder_outputs[1:]
  
          return BaseModelOutputWithPooling(
              last_hidden_state=sequence_output,
              pooler_output=pooled_output,
              hidden_states=encoder_outputs.hidden_states,
              attentions=encoder_outputs.attentions,
          )
  
      return pp_forward
  
file_path: ColossalAI/colossalai/shardformer/modeling/vit.py
imports: import math
  from typing import List, Optional, Tuple, Union
  import torch
  from transformers.models.vit.modeling_vit import BaseModelOutput, ViTEncoder
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
label_desc: Use image node_id 2 for calling a global function with example usage: ViTModel_pipeline_forward(stage_manager, stage_index) and returns: pp_forward, BaseModelOutputWithPooling, dict, unknown