code: def masked_language_model(
      *,
      X,
      mlm_weights,
      mlm_ids,
      mlm_positions,
      embed_weights,
      hidden,
      config,
      reuse=None,
      train=False,
  ):
      with tf.compat.v1.variable_scope("model/masked-language-model"):
          batch, seq, feats = shape_list(hidden)
          flat_offsets = tf.reshape(
              tf.range(0, batch, dtype=tf.int32) * seq, [-1, 1]
          )
  
          not_padding = tf.reshape(mlm_weights, [-1]) > 1e-9
          flat_positions = tf.boolean_mask(
              tensor=tf.reshape(mlm_positions + flat_offsets, [-1]),
              mask=not_padding,
          )  # take off the padding entirely
          gathered_hidden = tf.gather(
              tf.reshape(hidden, [batch * seq, feats]), flat_positions
          )
          mlm_ids = tf.boolean_mask(
              tensor=tf.reshape(mlm_ids, [-1]), mask=not_padding
          )
  
          final_proj_w = tf.compat.v1.get_variable(
              "dense/kernel",
              [config.n_embed, config.n_embed],
              initializer=tf.compat.v1.random_normal_initializer(
                  stddev=config.weight_stddev
              ),
          )
          final_proj_b = tf.compat.v1.get_variable(
              "dense/bias",
              [config.n_embed],
              initializer=tf.compat.v1.zeros_initializer,
          )
          final_proj = act_fns[config.act_fn](
              tf.matmul(gathered_hidden, final_proj_w, transpose_b=True)
              + final_proj_b
          )
  
          normed_proj = norm(final_proj, "LayerNorm")
          n_vocab = shape_list(embed_weights)[0]
          output_bias = tf.compat.v1.get_variable(
              "output_bias",
              shape=[n_vocab],
              initializer=tf.compat.v1.zeros_initializer(),
          )
  
          logits = tf.matmul(
              normed_proj, embed_weights, transpose_b=True
          )
          logits = tf.nn.bias_add(logits, output_bias)
  
          mlm_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
              logits=logits,
              labels=mlm_ids,
          )  # No weights needed as there is no padding.
          logits = tf.scatter_nd(
              indices=flat_positions,
              updates=logits,
              shape=[batch * seq, n_vocab],
          )
  
          return {
              "logits": logits,
              "losses": mlm_loss,
          }
  
imports: import functools
  import tensorflow
  from tensorflow_addons.text.crf import crf_log_likelihood
  from finetune.base_models.gpt.featurizer import attn, dropout, norm
  from finetune.util.shapes import shape_list, merge_leading_dims
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.errors import FinetuneError
  from finetune.nn.activations import act_fns
  from finetune.nn.nn_utils import norm
  from tensorflow.python.framework import function
label_desc: Use image node_id 2 for calling a global function with example usage: masked_language_model() and returns: dict