class_docstr: Torchscriptable version of `BertSelfAttention` from Huggingface transformers v2.3.0
  https://github.com/huggingface/transformers/blob/v2.3.0/transformers/modeling_bert.py # noqa
  
  Modifies the `forward` function and `transpose_for_scores` function
  
  Changes to `transpose_for_scores` function ::
      Changes the `new_x_shape` unpacking as static size inference is not supported
  
  Changes to `forward` function ::
      Uses scriptable `nn.functional.softmax` and also removes several static size
      inference which is not supported.
class_docstr_tok: ["Torchscriptable","version","of","`","BertSelfAttention","`","from","Huggingface","transformers","v2.3.0","https",":","\/\/github.com\/huggingface\/transformers\/blob\/v2.3.0\/transformers\/modeling_bert.py","#","noqa","Modifies","the","`","forward","`","function","and","`","transpose_for_scores","`","function","Changes","to","`","transpose_for_scores","`","function",":",":","Changes","the","`","new_x_shape","`","unpacking","as","static","size","inference","is","not","supported","Changes","to","`","forward","`","function",":",":","Uses","scriptable","`","nn.functional.softmax","`","and","also","removes","several","static","size","inference","which","is","not","supported","."]
code: def forward(
      self,
      hidden_states: Tensor,
      attention_mask: Optional[Tensor] = None,
      head_mask: Optional[Tensor] = None,
      encoder_hidden_states: Optional[Tensor] = None,
      encoder_attention_mask: Optional[Tensor] = None,
  ) -> Tuple[Tensor, Tensor]:
      mixed_query_layer = self.query(hidden_states)
  
      # If this is instantiated as a cross-attention module, the keys
      # and values come from an encoder; the attention mask needs to be
      # such that the encoder's padding tokens are not attended to.
      if encoder_hidden_states is not None:
          mixed_key_layer = self.key(encoder_hidden_states)
          mixed_value_layer = self.value(encoder_hidden_states)
          attention_mask = encoder_attention_mask
      else:
          mixed_key_layer = self.key(hidden_states)
          mixed_value_layer = self.value(hidden_states)
  
      query_layer = self.transpose_for_scores(mixed_query_layer)
      key_layer = self.transpose_for_scores(mixed_key_layer)
      value_layer = self.transpose_for_scores(mixed_value_layer)
  
      # Take the dot product between "query" and "key" to get the raw attention
      # scores.
      attention_scores = torch.matmul(
          query_layer, key_layer.transpose(-1, -2)
      )
      attention_scores = attention_scores / math.sqrt(
          self.attention_head_size
      )
      if attention_mask is not None:
          # Apply the attention mask is (precomputed for all layers in BertModel
          # forward() function)
          attention_scores = attention_scores + attention_mask
  
      # Normalize the attention scores to probabilities.
      attention_probs = nn.functional.softmax(attention_scores, dim=-1)
  
      # This is actually dropping out entire tokens to attend to, which might
      # seem a bit unusual, but is taken from the original Transformer paper.
      attention_probs = self.dropout(attention_probs)
  
      # Mask heads if we want to
      if head_mask is not None:
          attention_probs = attention_probs * head_mask
  
      context_layer = torch.matmul(attention_probs, value_layer)
  
      context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
      new_context_layer_shape = context_layer.size()[:-2] + (
          self.all_head_size,
      )
      context_layer = context_layer.view(new_context_layer_shape)
  
      outputs = (context_layer, attention_probs)
      return outputs
  
imports: import math
  from typing import List, Optional, Tuple
  import torch
  from mmf.utils.patch import restore_saved_modules, safecopy_modules
  from torch import nn, Tensor
label_desc: Use image node_id 2 for calling the BertSelfAttentionJit obj's underlying member method code with example usage: obj.forward(hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask) and returns: outputs