def _build_transformer_decoder(encoder_output,
                               source_len,
                               decoder_input,
                               mode,
                               model_config,
                               single_step_index=None):
  """Generate decoder output given encoder output and decoder input.

  Args:
    encoder_output: Tensor of shape (B, L, D).
    source_len: Tensor of shape (B) used to mask encoder_output.
    decoder_input: Tensor of shape (B, L, D).
    mode: Enum indicating model mode.
    model_config: ModelConfig proto.
    single_step_index: Optional scalar integer Tensor specifying an index in
      decoder_input. If set, then decoder will only calculate output for a
      single step. This can only be used during inference and if using a
      single-layer decoder, but significantly speeds up auto-regresssive
      decoding. Index should be >= 0 and < decoder_input length.

  Returns:
    decoder_output: the decoder output Tensor of shape (B, L, D)

  Raises:
    ValueError: Invalid inputs for single step decoding.
  """
  _set_initializer()

  # Only apply dropout during training, not eval or inference.
  perform_dropout = (mode == tf_estimator.ModeKeys.TRAIN)
  if perform_dropout:
    decoder_input = tf.nn.dropout(
        decoder_input,
        keep_prob=1.0 - model_config.training_options.layer_dropout_rate)

  # Get batch shape information based on input and output Tensors.
  batch_size = tf.shape(decoder_input)[0]

  # Attention bias vectors ensure padding in input sequence is ignored.
  max_encoder_length = tf.shape(encoder_output)[1]
  encoder_decoder_attention_bias = common_layers.attention_bias_ignore_padding(
      source_len, max_length=max_encoder_length)

  # This attention bias matrix ensures decoder only has access to past
  # symbols.
  max_decoder_length = tf.shape(decoder_input)[1]
  decoder_self_attention_bias = common_layers.attention_bias_lower_triangle(
      max_decoder_length)

  # Setup decoder adjacency matrix.
  adjacency_matrix = common_layers.relative_positions_adjacency_matrix(
      max_decoder_length, batch_size,
      model_config.model_parameters.max_decoder_relative_distance)
  num_labels = 2 * model_config.model_parameters.max_decoder_relative_distance + 1

  x = decoder_input
  memory_antecedent = decoder_input

  using_single_step_optimization = False

  # If set, only decode a single step, otherwise decode all steps in parallel.
  if single_step_index is not None:
    if mode != tf_estimator.ModeKeys.PREDICT:
      raise ValueError("Single step decoding only supported during inference.")

    # Important optimization if using a single layer decoder.
    # Instead of running a masked decoder self-attention over all decoder
    # inputs, only run attention of current the step against all other decoder
    # inputs.
    # Need to ensure we are only using a single layer decoder.
    # Since single step decoding only computes a single decoder step,
    # higher layers would have incomplete information for self attention
    # with this optimization.
    if model_config.model_parameters.num_decoder_layers == 1:
      using_single_step_optimization = True
      x = tf.gather(decoder_input, [single_step_index], axis=1)
      adjacency_matrix = tf.gather(
          adjacency_matrix, [single_step_index], axis=1)
      decoder_self_attention_bias = tf.gather(decoder_self_attention_bias,
                                              [single_step_index])

  # Stack of decoder layers.
  with tf.variable_scope("decoder"):
    for layer in range(model_config.model_parameters.num_decoder_layers):
      with tf.variable_scope("layer_%d" % layer):
        with tf.variable_scope("self_attention"):
          y = common_layers.graph_attention(
              x,
              memory_antecedent,
              decoder_self_attention_bias,
              model_config.model_parameters.decoder_dims,
              model_config.model_parameters.num_heads,
              adjacency_matrix=adjacency_matrix,
              num_labels=num_labels)
          x = common_layers.layer_postprocess(x, y, model_config,
                                              perform_dropout)
        with tf.variable_scope("encdec_attention"):
          y = common_layers.attention(
              x, encoder_output, encoder_decoder_attention_bias,
              model_config.model_parameters.decoder_dims,
              model_config.model_parameters.num_heads)
          x = common_layers.layer_postprocess(x, y, model_config,
                                              perform_dropout)
        with tf.variable_scope("ffn"):
          y = common_layers.ff_layer(
              x, model_config.model_parameters.decoder_ff_layer_hidden_size,
              model_config.model_parameters.decoder_dims)
          x = common_layers.layer_postprocess(x, y, model_config,
                                              perform_dropout)
          memory_antecedent = x
  if single_step_index is not None and not using_single_step_optimization:
    x = tf.gather(x, [single_step_index], axis=1)
  return x
Generate decoder output given encoder output and decoder input.

  Args:
    encoder_output: Tensor of shape (B, L, D).
    source_len: Tensor of shape (B) used to mask encoder_output.
    decoder_input: Tensor of shape (B, L, D).
    mode: Enum indicating model mode.
    model_config: ModelConfig proto.
    single_step_index: Optional scalar integer Tensor specifying an index in
      decoder_input. If set, then decoder will only calculate output for a
      single step. This can only be used during inference and if using a
      single-layer decoder, but significantly speeds up auto-regresssive
      decoding. Index should be >= 0 and < decoder_input length.

  Returns:
    decoder_output: the decoder output Tensor of shape (B, L, D)

  Raises:
    ValueError: Invalid inputs for single step decoding.
  
language/xsp/model/transformer.py