class_docstr: Top-k sequence generator based on the decoder followed by log_softmax.
  
  Args:
      *all args of GreedySequenceGenerator class
      beam_size: size of the beam (parameter k in top-k)
      temperature: temperature of top-k sampling, all logits are divided
          by temperature before rescaling. High temperature leads to
          uniform distribution, low leads to delta-like distribution.
  Kwargs:
      all remaining parameters of GreedySequenceGenerator class
class_docstr_tok: ["Top-k","sequence","generator","based","on","the","decoder","followed","by","log_softmax",".","Args",":","*","all","args","of","GreedySequenceGenerator","class","beam_size",":","size","of","the","beam","(","parameter","k","in","top-k",")","temperature",":","temperature","of","top-k","sampling",",","all","logits","are","divided","by","temperature","before","rescaling",".","High","temperature","leads","to","uniform","distribution",",","low","leads","to","delta-like","distribution",".","Kwargs",":","all","remaining","parameters","of","GreedySequenceGenerator","class"]
code: def _one_step_forward(
      self,
      decoder_input_ids=None,
      encoder_hidden_states=None,
      encoder_input_mask=None,
      decoder_mems_list=None,
      pos=0,
  ):
      log_probs, decoder_mems_list = super()._one_step_forward(
          decoder_input_ids,
          encoder_hidden_states,
          encoder_input_mask,
          decoder_mems_list,
          pos,
      )
  
      batch_size, seq_len, vocab_size = log_probs.size()
      scores, indices = torch.topk(log_probs, self.beam_size, dim=-1)
  
      rescaled_logexp = torch.zeros_like(log_probs).scatter(
          -1, indices, scores.div(self.temp).exp()
      )
      probs = rescaled_logexp / rescaled_logexp.norm(
          1, -1, keepdim=True
      )
  
      # We randomly sample next tokens from rescaled probability distribution
      # over top-k candidates and return a binary tensor which indicates
      # candidates that have been selected. We call this object
      # `pseudo_log_probs` as genuine log_probs should have -infs instead of
      # 0s and 0s instead of 1s.
      ids = torch.multinomial(probs.view(-1, vocab_size), 1).view(
          -1, seq_len, 1
      )
      pseudo_log_probs = torch.zeros_like(log_probs).scatter(
          -1, ids, 1.0
      )
  
      return pseudo_log_probs, decoder_mems_list
  
imports: from contextlib import contextmanager
  import torch
  from nemo.collections.common.parts import NEG_INF, mask_padded_tokens
label_desc: Use image node_id 2 for calling the TopKSequenceGenerator obj's underlying member method code with example usage: obj._one_step_forward(decoder_input_ids, encoder_hidden_states, encoder_input_mask, decoder_mems_list, pos) and returns: pseudo_log_probs,  decoder_mems_list