docstr: If masked_tokens_mask is not None (i.e. last_layer_subset == True in BertForPreTraining),
  we only want the output for the masked tokens. This means that we only compute the last
  layer output for these tokens.
  masked_tokens_mask: (batch, seqlen), dtype=torch.bool
docstr_tok: ["If","masked_tokens_mask","is","not","None","(","i.e",".","last_layer_subset","==","True","in","BertForPreTraining",")",",","we","only","want","the","output","for","the","masked","tokens",".","This","means","that","we","only","compute","the","last","layer","output","for","these","tokens",".","masked_tokens_mask",":","(","batch",",","seqlen",")",",","dtype=torch.bool"]
code: def forward(
      self,
      input_ids,
      position_ids=None,
      token_type_ids=None,
      attention_mask=None,
      masked_tokens_mask=None,
  ):
      """If masked_tokens_mask is not None (i.e. last_layer_subset == True in BertForPreTraining),
      we only want the output for the masked tokens. This means that we only compute the last
      layer output for these tokens.
      masked_tokens_mask: (batch, seqlen), dtype=torch.bool
      """
      hidden_states = self.embeddings(
          input_ids,
          position_ids=position_ids,
          token_type_ids=token_type_ids,
      )
      # TD [2022-12:18]: Don't need to force residual in fp32
      # BERT puts embedding LayerNorm before embedding dropout.
      if not self.fused_dropout_add_ln:
          hidden_states = self.emb_ln(hidden_states)
      else:
          hidden_states = layer_norm_fn(
              hidden_states,
              self.emb_ln.weight,
              self.emb_ln.bias,
              eps=self.emb_ln.eps,
          )
      hidden_states = self.emb_drop(hidden_states)
  
      if masked_tokens_mask is not None:
          batch_size, seqlen = input_ids.shape[:2]
          # We also need the first column for the CLS token
          first_col_mask = torch.zeros(
              batch_size,
              seqlen,
              dtype=torch.bool,
              device=input_ids.device,
          )
          first_col_mask[:, 0] = True
          subset_mask = masked_tokens_mask | first_col_mask
      else:
          subset_mask = None
  
      sequence_output = self.encoder(
          hidden_states,
          key_padding_mask=attention_mask,
          subset_mask=subset_mask,
      )
  
      if masked_tokens_mask is None:
          pooled_output = (
              self.pooler(sequence_output)
              if self.pooler is not None
              else None
          )
      else:
          # TD [2022-03-01]: the indexing here is very tricky.
          if attention_mask is not None:
              subset_idx = subset_mask[attention_mask]
              pool_input = sequence_output[
                  first_col_mask[attention_mask][subset_idx]
              ]
              sequence_output = sequence_output[
                  masked_tokens_mask[attention_mask][subset_idx]
              ]
          else:
              pool_input = sequence_output[first_col_mask[subset_mask]]
              sequence_output = sequence_output[
                  masked_tokens_mask[subset_mask]
              ]
          pooled_output = (
              self.pooler(pool_input, pool=False)
              if self.pooler is not None
              else None
          )
  
      return BaseModelOutputWithPoolingAndCrossAttentions(
          last_hidden_state=sequence_output,
          pooler_output=pooled_output,
      )
  
imports: import logging
  import re
  from collections import OrderedDict
  from collections.abc import Sequence
  from functools import partial
  from typing import Any, Mapping
  import torch
  import torch.nn
  import torch.nn.functional
  from einops import rearrange
  from transformers import BertConfig, PretrainedConfig
  from transformers.models.bert.modeling_bert import BaseModelOutputWithPoolingAndCrossAttentions, BertForPreTrainingOutput
  from flash_attn.bert_padding import index_first_axis, index_first_axis_residual, pad_input, unpad_input
  from flash_attn.modules.block import Block
  from flash_attn.modules.embedding import BertEmbeddings
  from flash_attn.modules.mha import MHA
  from flash_attn.modules.mlp import FusedMLP, Mlp
  from flash_attn.utils.pretrained import state_dict_from_pretrained
label_desc: Use image node_id 2 for calling the BertModel obj's underlying member method code with example usage: obj.forward(input_ids, position_ids, token_type_ids, attention_mask, masked_tokens_mask) and returns: BaseModelOutputWithPoolingAndCrossAttentions