class_docstr: Masked language model network head for BERT modeling.
  
  This network implements a masked language model based on the provided network.
  It assumes that the network being passed has a "get_embedding_table()" method.
  
  Attributes:
    input_width: The innermost dimension of the input tensor to this network.
    num_predictions: The number of predictions to make per sequence.
    source_network: The network with the embedding layer to use for the
      embedding layer.
    activation: The activation, if any, for the dense layer in this network.
    initializer: The intializer for the dense layer in this network. Defaults to
      a Glorot uniform initializer.
    output: The output style for this network. Can be either 'logits' or
      'predictions'.
class_docstr_tok: ["Masked","language","model","network","head","for","BERT","modeling",".","This","network","implements","a","masked","language","model","based","on","the","provided","network",".","It","assumes","that","the","network","being","passed","has","a","``","get_embedding_table","(",")","''","method",".","Attributes",":","input_width",":","The","innermost","dimension","of","the","input","tensor","to","this","network",".","num_predictions",":","The","number","of","predictions","to","make","per","sequence",".","source_network",":","The","network","with","the","embedding","layer","to","use","for","the","embedding","layer",".","activation",":","The","activation",",","if","any",",","for","the","dense","layer","in","this","network",".","initializer",":","The","intializer","for","the","dense","layer","in","this","network",".","Defaults","to","a","Glorot","uniform","initializer",".","output",":","The","output","style","for","this","network",".","Can","be","either","'logits","'","or","'predictions","'","."]
code: def __init__(
      self,
      input_width,
      num_predictions,
      source_network,
      activation=None,
      initializer="glorot_uniform",
      output="logits",
      **kwargs,
  ):
      embedding_table = source_network.get_embedding_table()
      vocab_size, hidden_size = embedding_table.shape
  
      sequence_data = tf.keras.layers.Input(
          shape=(None, input_width),
          name="sequence_data",
          dtype=tf.float32,
      )
      masked_lm_positions = tf.keras.layers.Input(
          shape=(num_predictions,),
          name="masked_lm_positions",
          dtype=tf.int32,
      )
  
      masked_lm_input = tf.keras.layers.Lambda(
          lambda x: self._gather_indexes(x[0], x[1])
      )([sequence_data, masked_lm_positions])
      lm_data = tf.keras.layers.Dense(
          hidden_size,
          activation=activation,
          kernel_initializer=initializer,
          name="cls/predictions/transform/dense",
      )(masked_lm_input)
      lm_data = tf.keras.layers.LayerNormalization(
          axis=-1,
          epsilon=1e-12,
          name="cls/predictions/transform/LayerNorm",
      )(lm_data)
      lm_data = tf.keras.layers.Lambda(
          lambda x: tf.matmul(x, embedding_table, transpose_b=True)
      )(lm_data)
      # lm_data = tf.linalg.matmul(lm_data, embedding_table, transpose_b=True)
      logits = Bias(
          initializer=tf.keras.initializers.Zeros(),
          name="cls/predictions/output_bias",
      )(lm_data)
  
      # We can't use the standard Keras reshape layer here, since it expects
      # the input and output batch size to be the same.
      reshape_layer = tf.keras.layers.Lambda(
          lambda x: tf.reshape(x, [-1, num_predictions, vocab_size])
      )
  
      self.logits = reshape_layer(logits)
      predictions = tf.keras.layers.Activation(tf.nn.log_softmax)(
          self.logits
      )
  
      if output == "logits":
          output_tensors = self.logits
      elif output == "predictions":
          output_tensors = predictions
      else:
          raise ValueError(
              (
                  'Unknown `output` value "%s". `output` can be either "logits" or '
                  '"predictions"'
              )
              % output
          )
  
      super(MaskedLM, self).__init__(
          inputs=[sequence_data, masked_lm_positions],
          outputs=output_tensors,
          **kwargs,
      )
  
file_path: autodist/examples/benchmark/utils/modeling/networks/masked_lm.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tensorflow.python.keras.engine import network
  from utils import tf_utils
label_desc: Use image node_id 1 to create a new MaskedLM object from inherited base classes: network with example: obj = MaskedLM(input_width, num_predictions, source_network, activation, initializer, output)