class_docstr: This class serves as a micro library for forward function substitution of Bert models
  under pipeline setting.
docstr: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
      Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
      config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
      loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
docstr_tok: ["labels","(","`","torch.LongTensor","`","of","shape","`","(","batch_size",",","sequence_length",")","`",",","*","optional","*",")",":","Labels","for","computing","the","masked","language","modeling","loss",".","Indices","should","be","in","`","[","-100",",","0",",","...",",","config.vocab_size","]","`","(","see","`","input_ids","`","docstring",")","Tokens","with","indices","set","to","`","-100","`","are","ignored","(","masked",")",",","the","loss","is","only","computed","for","the","tokens","with","labels","in","`","[","0",",","...",",","config.vocab_size","]","`"]
code: def bert_for_masked_lm_forward(
      self: BertForMaskedLM,
      input_ids: Optional[torch.Tensor] = None,
      attention_mask: Optional[torch.Tensor] = None,
      token_type_ids: Optional[torch.Tensor] = None,
      position_ids: Optional[torch.Tensor] = None,
      head_mask: Optional[torch.Tensor] = None,
      inputs_embeds: Optional[torch.Tensor] = None,
      encoder_hidden_states: Optional[torch.Tensor] = None,
      encoder_attention_mask: Optional[torch.Tensor] = None,
      labels: Optional[torch.Tensor] = None,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
      hidden_states: Optional[torch.Tensor] = None,
      stage_manager: Optional[PipelineStageManager] = None,
      stage_index: Optional[List[int]] = None,
      shard_config: ShardConfig = None,
  ):
      r"""
      labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
          Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
          config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
          loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
      """
      logger = logging.get_logger(__name__)
  
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.use_return_dict
      )
  
      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
  
      outputs = BertPipelineForwards.bert_model_forward(
          self.bert,
          input_ids,
          attention_mask=attention_mask,
          token_type_ids=token_type_ids,
          position_ids=position_ids,
          head_mask=head_mask,
          inputs_embeds=inputs_embeds,
          encoder_hidden_states=encoder_hidden_states,
          encoder_attention_mask=encoder_attention_mask,
          output_attentions=output_attentions,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
          hidden_states=hidden_states,
          stage_manager=stage_manager,
          stage_index=stage_index,
          shard_config=shard_config,
      )
  
      if stage_manager.is_last_stage():
          sequence_output = outputs[0]
          prediction_scores = self.cls(sequence_output)
  
          masked_lm_loss = None
          if labels is not None:
              loss_fct = (
                  CrossEntropyLoss()
              )  # -100 index = padding token
              masked_lm_loss = loss_fct(
                  prediction_scores.view(-1, self.config.vocab_size),
                  labels.view(-1),
              )
  
          if not return_dict:
              output = (prediction_scores,) + outputs[2:]
              return (
                  ((masked_lm_loss,) + output)
                  if masked_lm_loss is not None
                  else output
              )
  
          return MaskedLMOutput(
              loss=masked_lm_loss,
              logits=prediction_scores,
              hidden_states=outputs.hidden_states,
              attentions=outputs.attentions,
          )
      else:
          hidden_states = outputs.get("hidden_states")
          return {"hidden_states": hidden_states}
  
imports: import math
  import warnings
  from typing import List, Optional, Tuple, Union
  import torch
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions, CausalLMOutputWithCrossAttentions, MaskedLMOutput, MultipleChoiceModelOutput, NextSentencePredictorOutput, QuestionAnsweringModelOutput, SequenceClassifierOutput, TokenClassifierOutput
  from transformers.models.bert.modeling_bert import BertForMaskedLM, BertForMultipleChoice, BertForNextSentencePrediction, BertForPreTraining, BertForPreTrainingOutput, BertForQuestionAnswering, BertForSequenceClassification, BertForTokenClassification, BertLMHeadModel, BertModel
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from colossalai.shardformer import ShardConfig
  from colossalai.shardformer.layer._operation import gather_forward_split_backward, split_forward_gather_backward
label_desc: Use image node_id 4 for calling the BertPipelineForwards obj's underlying member method code with example usage: obj.bert_for_masked_lm_forward(input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, output_attentions, output_hidden_states, return_dict, hidden_states, stage_manager, stage_index, shard_config) and returns: MaskedLMOutput, dict, unknown