def train(model_config,
          input_embeddings,
          source_len,
          output_vocab_size,
          output_vocab_embeddings_table,
          target_decode_steps,
          mode,
          input_copy_mask=None):
  """Constructs encoder and decoder transformation for training and eval.

  In the shapes described below, B is batch size, L is sequence length,
  D is the dimensionality of the model embeddings, and T is the output vocab
  size.

  Args:
    model_config: ModelConfig proto.
    input_embeddings: Tensor of shape (B, L, D) representing inputs.
    source_len: Tensor of shape (B) containing length of each input sequence.
    output_vocab_size: Size of output vocabulary.
    output_vocab_embeddings_table: Tensor of shape (T, D) representing table of
      embeddings for output symbols.
    target_decode_steps: DecodeSteps representing target outputs. Each tensor
      has shape (B, L).
    mode: Enum indicating model mode, TRAIN or EVAL.
    input_copy_mask: Mask for copying actions.

  Returns:
    Tuple of (logits, predicted_ids), where logits is a tensor of shape
    (B, L, T) representing model output logits, and predicted_ids is
    a tensor of shape (B, L) containing the derived integer IDs of the
    one-best output symbol.
  """
  logits = _transformer_body(
      input_embeddings,
      source_len,
      target_decode_steps,
      mode,
      model_config,
      output_vocab_size,
      output_vocab_embeddings_table,
      input_copy_mask=input_copy_mask)

  predicted_ids = tf.to_int32(tf.argmax(logits, axis=-1))
  output_decode_steps = decode_utils.get_decode_steps(predicted_ids,
                                                      output_vocab_size,
                                                      model_config)
  predictions = decode_utils.get_predictions(output_decode_steps)
  return logits, predictions
Constructs encoder and decoder transformation for training and eval.

  In the shapes described below, B is batch size, L is sequence length,
  D is the dimensionality of the model embeddings, and T is the output vocab
  size.

  Args:
    model_config: ModelConfig proto.
    input_embeddings: Tensor of shape (B, L, D) representing inputs.
    source_len: Tensor of shape (B) containing length of each input sequence.
    output_vocab_size: Size of output vocabulary.
    output_vocab_embeddings_table: Tensor of shape (T, D) representing table of
      embeddings for output symbols.
    target_decode_steps: DecodeSteps representing target outputs. Each tensor
      has shape (B, L).
    mode: Enum indicating model mode, TRAIN or EVAL.
    input_copy_mask: Mask for copying actions.

  Returns:
    Tuple of (logits, predicted_ids), where logits is a tensor of shape
    (B, L, T) representing model output logits, and predicted_ids is
    a tensor of shape (B, L) containing the derived integer IDs of the
    one-best output symbol.
  
language/xsp/model/transformer.py