def _ComputeAttention(self,
                        theta,
                        rnn_out,
                        packed_src,
                        attention_state,
                        per_step_src_padding=None,
                        query_segment_id=None):
    """Runs attention and computes context vector.

    Can be overridden by a child class if attention is computed differently.

    Args:
      theta: A NestedMap object containing weights for the attention layers.
        Expects a member named 'atten'.
      rnn_out: A Tensor of shape [batch_size, query_dim]; output of the
        first layer of decoder RNN, which is the query vector used for
        attention.
      packed_src: A NestedMap returned by self.atten.InitForSourcePacked.
      attention_state: The attention state computed at the previous timestep.
        Varies with the type of attention, but is usually a Tensor or a
        NestedMap of Tensors of shape [batch_size, <state_dim>].
      per_step_src_padding: Source sequence padding to apply at this step.
      query_segment_id: a tensor of shape [batch_size].

    Returns:
      A tuple of 3 tensors:

      - The attention context vector: shaped [batch_size, context_dim].
      - The attention probability vector: shaped [batch_size, seq_len]
      - The attention state: A Tensor or a NestedMap of Tensors of shape
        [batch_size, <state_dim>].
    """
    return self.atten.ComputeContextVectorWithSource(
        theta.atten,
        packed_src,
        rnn_out,
        attention_state=attention_state,
        per_step_source_padding=per_step_src_padding,
        query_segment_id=query_segment_id)
Runs attention and computes context vector.

    Can be overridden by a child class if attention is computed differently.

    Args:
      theta: A NestedMap object containing weights for the attention layers.
        Expects a member named 'atten'.
      rnn_out: A Tensor of shape [batch_size, query_dim]; output of the
        first layer of decoder RNN, which is the query vector used for
        attention.
      packed_src: A NestedMap returned by self.atten.InitForSourcePacked.
      attention_state: The attention state computed at the previous timestep.
        Varies with the type of attention, but is usually a Tensor or a
        NestedMap of Tensors of shape [batch_size, <state_dim>].
      per_step_src_padding: Source sequence padding to apply at this step.
      query_segment_id: a tensor of shape [batch_size].

    Returns:
      A tuple of 3 tensors:

      - The attention context vector: shaped [batch_size, context_dim].
      - The attention probability vector: shaped [batch_size, seq_len]
      - The attention state: A Tensor or a NestedMap of Tensors of shape
        [batch_size, <state_dim>].
    
contrib/lingvo-patched-decoder.py