code: def step(
      self,
      step: int,
      lprobs,
      scores,
      prev_output_tokens: Optional[Tensor] = None,
      original_batch_idxs: Optional[Tensor] = None,
  ):
      bsz, beam_size, vocab_size = lprobs.size()
  
      if step == 0:
          # at the first step all hypotheses are equally likely, so use
          # only the first beam
          lprobs = lprobs[:, ::beam_size, :].contiguous()
  
      if self.sampling_topp > 0:
          # only sample from the smallest set of words whose cumulative probability mass exceeds p
          probs, top_indices = self._sample_topp(lprobs)
      elif self.sampling_topk > 0:
          # only sample from top-k candidates
          lprobs, top_indices = lprobs.topk(self.sampling_topk)
          probs = lprobs.exp_()
      else:
          probs = lprobs.exp_()
  
          # dummy data to be consistent with true branch for type check
          top_indices = torch.empty(0).to(probs)
      # sample
      if step == 0:
          indices_buf = torch.multinomial(
              probs.view(bsz, -1),
              beam_size,
              replacement=True,
          ).view(bsz, beam_size)
      else:
          indices_buf = torch.multinomial(
              probs.view(bsz * beam_size, -1),
              1,
              replacement=True,
          ).view(bsz, beam_size)
  
      if step == 0:
          # expand to beam size
          probs = probs.expand(bsz, beam_size, -1)
  
      # gather scores
      scores_buf = torch.gather(
          probs, dim=2, index=indices_buf.unsqueeze(-1)
      )
      scores_buf = scores_buf.log_().view(bsz, -1)
  
      # remap indices if using top-k or top-P sampling
      if self.sampling_topk > 0 or self.sampling_topp > 0:
          indices_buf = torch.gather(
              top_indices.expand(bsz, beam_size, -1),
              dim=2,
              index=indices_buf.unsqueeze(-1),
          ).squeeze(2)
  
      if step == 0:
          beams_buf = indices_buf.new_zeros(bsz, beam_size)
      else:
          beams_buf = (
              torch.arange(0, beam_size).to(indices_buf).repeat(bsz, 1)
          )
          # make scores cumulative
          scores_buf.add_(
              torch.gather(
                  scores[:, :, step - 1], dim=1, index=beams_buf
              )
          )
  
      return scores_buf, indices_buf, beams_buf
  
imports: import math
  from typing import List, Optional
  import torch
  import torch.nn
  from fairseq.token_generation_constraints import ConstraintState, OrderedConstraintState, UnorderedConstraintState
  from torch import Tensor
label_desc: Use image node_id 3 for calling the Sampling obj's underlying member method code with example usage: obj.step(step, lprobs, scores, prev_output_tokens, original_batch_idxs) and returns: scores_buf,  indices_buf,  beams_buf