def PostStepDecoderStateUpdate(self, decoder_step_state, logits=None):
    """Update decoder states and logits after SingleDecodeStep.

    Args:
      decoder_step_state: A NestedMap object which encapsulates decoder states.
      logits: a tensor, predicted logits.

    Returns:
      decoder_step_state.

    Raises:
      ValueError: if scheduled sampling is used for functional unrolling or
                  if logits is None for while loop based unrolling.
    """
    if not self.params.use_while_loop_based_unrolling:
      if self.params.min_ground_truth_prob < 1.0:
        raise ValueError('SS is not yet supported')
    else:
      if logits is None:
        raise ValueError('logits cannot be None')
      decoder_step_state.logits = logits

      if self._max_label_prob > 0:
        bs = tf.shape(logits)[0]
        # log_probs: [bs, num_classes]
        log_probs = tf.nn.log_softmax(logits)
        # log_prob_sample: [bs, 1]
        log_prob_sample = tf.multinomial(
            log_probs, 1, seed=self.params.random_seed)
        # pred_ids: [bs]
        pred_ids = tf.reshape(tf.cast(log_prob_sample, tf.int32), [bs])
        decoder_step_state.misc_states.prev_predicted_ids = pred_ids
    return decoder_step_state
Update decoder states and logits after SingleDecodeStep.

    Args:
      decoder_step_state: A NestedMap object which encapsulates decoder states.
      logits: a tensor, predicted logits.

    Returns:
      decoder_step_state.

    Raises:
      ValueError: if scheduled sampling is used for functional unrolling or
                  if logits is None for while loop based unrolling.
    
contrib/lingvo-patched-decoder.py