class_docstr: BERT model ("Bidirectional Encoder Representations from Transformers").
class_docstr_tok: ["BERT","model","(","``","Bidirectional","Encoder","Representations","from","Transformers","''",")","."]
docstr: Performs multi-headed self-attention with output projection.
  
  Args:
    x: output of previous layer
    attention_bias: optional float32 Tensor broadcastable to shape
      x.shape - self.model_dim + self.memory_seq_dim
      to be added to attention logits.
      This may used to mask out padding regions of the memory.
  
  Returns:
    float Tensor with the same shape as x
code: def self_attention(self, x, attention_bias):
      """Performs multi-headed self-attention with output projection.
  
      Args:
        x: output of previous layer
        attention_bias: optional float32 Tensor broadcastable to shape
          x.shape - self.model_dim + self.memory_seq_dim
          to be added to attention logits.
          This may used to mask out padding regions of the memory.
  
      Returns:
        float Tensor with the same shape as x
      """
  
      queries = mtf.layers.dense(
          x,
          reduced_dims=[self.model_dim],
          new_dims=[self.num_heads_dim, self.size_per_head_dim],
          kernel_initializer=self.dense_initializer,
          name="query",
          use_bias=self.config.use_bias,
      )
      keys = mtf.layers.dense(
          mtf.replace_dimensions(x, self.seq_dim, self.memory_seq_dim),
          reduced_dims=[self.model_dim],
          new_dims=[self.num_heads_dim, self.size_per_head_dim],
          kernel_initializer=self.dense_initializer,
          name="key",
          use_bias=self.config.use_bias,
      )
      values = mtf.layers.dense(
          mtf.replace_dimensions(x, self.seq_dim, self.memory_seq_dim),
          reduced_dims=[self.model_dim],
          new_dims=[self.num_heads_dim, self.size_per_head_dim],
          kernel_initializer=self.dense_initializer,
          name="value",
          use_bias=self.config.use_bias,
      )
  
      # Take the dot product between "query" and "key" to get the raw
      # attention scores.
      attention_scores = mtf.einsum(
          [queries, keys], reduced_dims=[self.size_per_head_dim]
      )
      attention_scores *= self.size_per_head_dim.size**-0.5
  
      if attention_bias is not None:
          attention_scores += attention_bias
  
      # Normalize the attention scores to probabilities.
      attention_probs = mtf.softmax(
          attention_scores, self.memory_seq_dim
      )
  
      # 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 = mtf.dropout(
          attention_probs,
          is_training=(self.config.attention_probs_dropout_prob == 0.0),
          keep_prob=1.0 - self.config.attention_probs_dropout_prob,
      )
  
      output = mtf.einsum(
          [attention_probs, values], reduced_dims=[self.memory_seq_dim]
      )
  
      # linear transformation back to shape of query_antecedent
      output = mtf.layers.dense(
          output,
          reduced_dims=[self.num_heads_dim, self.size_per_head_dim],
          new_dims=[self.model_dim],
          kernel_initializer=self.dense_initializer,
          name="output",
          use_bias=self.config.use_bias,
      )
      output = mtf.transpose(output, x.shape)
      return output
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import copy
  import json
  import math
  import re
  import mesh_tensorflow
  import mesh_tensorflow.transformer.moe
  import six
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling the BertModel obj's underlying member method code with example usage: obj.self_attention(x, attention_bias) and returns: output