def transformer_error_tag_prediction_layer(x,
                                           hparams,
                                           features,
                                           loss_mask,
                                           layer_collection=None):
  """Layer that predicts the error tag."""
  with tf.variable_scope("error_tag_prediction"):
    x = maybe_flatten4d3d(x)
    vocab_size = hparams.problem.feature_info["targets_error_tag"].vocab_size
    labels = features["targets_error_tag_raw"]
    with tf.variable_scope("projection"):
      bottleneck = common_layers.dense(
          x,
          hparams.error_tag_embed_size,
          layer_collection=layer_collection,
          name="bottleneck")
      logits = common_layers.dense(
          bottleneck,
          vocab_size,
          use_bias=False,
          layer_collection=layer_collection,
          name="logits")
      xent = tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=logits, labels=labels)
      loss = tf.reduce_sum(xent * loss_mask)
    with tf.variable_scope("embedding"):
      embed_mat = get_error_tag_embedding_matrix()
      y = common_layers.layer_preprocess(
          common_layers.embedding(
              labels, vocab_size, hparams.hidden_size, embedding_var=embed_mat),
          hparams,
          layer_collection=layer_collection)
      x = common_layers.layer_postprocess(x, y, hparams)
    return x, logits, loss
Layer that predicts the error tag.
tensor2tensor/models/research/transformer_seq2edits.py