class_docstr: Step-by-step decoder with LM fusion.
class_docstr_tok: ["Step-by-step","decoder","with","LM","fusion","."]
code: def _PreBeamSearchStepCallback(
      self, theta, encoder_outputs, step_ids, states, num_hyps_per_beam
  ):
      p = self.params
      step_paddings = tf.zeros(tf.shape(step_ids), dtype=p.dtype)
      embs = self.emb.EmbLookup(theta.emb, tf.reshape(step_ids, [-1]))
      prev_rnn_states = states.rnn_states
      prev_atten_states = states.all_atten_states.atten_states
      prev_atten_context = states.all_atten_states.atten_context
      prev_atten_probs = states.all_atten_states.atten_probs
      prev_fusion_states = states.fusion_states
      prev_misc_states = states.misc_states
  
      prev_decoder_step_state = py_utils.NestedMap(
          rnn_states=prev_rnn_states,
          atten_context=prev_atten_context,
          atten_probs=prev_atten_probs,
          atten_states=prev_atten_states,
          misc_states=prev_misc_states,
      )
      # TODO(prabhavalkar): Must handle CreateMiscTargetInfo during beam search
      # eval.
      cur_target_info = AsrDecoderBase.TargetInfo(
          id=tf.reshape(step_ids, [-1]),
          label=None,
          weight=None,
          emb=embs,
          padding=step_paddings,
          misc=py_utils.NestedMap(),
      )
  
      packed_src = self._InitAttention(theta, encoder_outputs)
      step_out, new_decoder_step_state = self.SingleDecodeStep(
          theta,
          packed_src,
          cur_target_info=cur_target_info,
          decoder_step_state=prev_decoder_step_state,
      )
      (
          atten_context,
          atten_probs,
          rnn_states,
          atten_states,
          misc_states,
      ) = (
          new_decoder_step_state.atten_context,
          new_decoder_step_state.atten_probs,
          new_decoder_step_state.rnn_states,
          new_decoder_step_state.atten_states,
          new_decoder_step_state.misc_states,
      )
  
      if p.softmax_uses_attention:
          # [batch, dims]
          softmax_input = step_out
      else:
          # Strip the attention context from the last dimension of softmax_input.
          # TODO(prabhavalkar): This currently assumes that the context is appended
          # to the end, see tf.concat in
          # AsrDecoderBase.ComputePredictionsFunctional().RnnStep().  Refactor the
          # code so as to remove this assumption.
          atten_context_dim = self._GetAttenContextDim()
          rnn_output_dim = self.rnn_cell[-1].params.num_output_nodes
          softmax_input, _ = tf.split(
              step_out, [rnn_output_dim, atten_context_dim], axis=-1
          )
  
      softmax_input, fusion_states = self.fusion.FProp(
          theta.fusion,
          prev_fusion_states,
          softmax_input,
          cur_target_info.id,
          cur_target_info.padding,
      )
  
      logits = self._ComputeLogits(theta, softmax_input)
      logits = self.fusion.ComputeLogitsWithLM(
          fusion_states, logits, is_eval=True
      )
      if p.use_unnormalized_logits_as_log_probs:
          log_probs = logits
      else:
          log_probs = tf.nn.log_softmax(logits)
  
      atten_probs = self._PostProcessAttenProbsForBeamSearch(
          atten_probs
      )
      bs_results = py_utils.NestedMap(
          {
              "atten_probs": atten_probs,
              "log_probs": log_probs,
          }
      )
      all_atten_states = py_utils.NestedMap(
          {
              "atten_context": atten_context,
              "atten_probs": atten_probs,
              "atten_states": atten_states,
          }
      )
      new_states = py_utils.NestedMap(
          {
              "rnn_states": rnn_states,
              "all_atten_states": all_atten_states,
              "fusion_states": fusion_states,
              "misc_states": misc_states,
          }
      )
      return bs_results, new_states
  
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 8 for calling the AsrDecoder obj's underlying member method code with example usage: obj._PreBeamSearchStepCallback(theta, encoder_outputs, step_ids, states, num_hyps_per_beam) and returns: bs_results,  new_states