class_docstr: Step-by-step decoder with LM fusion.
class_docstr_tok: ["Step-by-step","decoder","with","LM","fusion","."]
docstr: 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>].
docstr_tok: ["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",">","]","."]
code: 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,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import math
  import lingvo.compat
  from lingvo.core import attention
  from lingvo.core import base_decoder
  from lingvo.core import base_layer
  from lingvo.core import cluster_factory
  from lingvo.core import layers
  from lingvo.core import plot
  from lingvo.core import py_utils
  from lingvo.core import recurrent
  from lingvo.core import rnn_cell
  from lingvo.core import summary_utils
  from lingvo.core import symbolic
  from lingvo.tasks.asr import contextualizer_base
  from lingvo.tasks.asr import decoder_utils
  from lingvo.tasks.asr import fusion
  from matplotlib import font_manager
  import six
  from six.moves import range
label_desc: Use image node_id 3 for calling the AsrDecoder obj's underlying member method code with example usage: obj._ComputeAttention(theta, rnn_out, packed_src, attention_state, per_step_src_padding, query_segment_id) and returns: self