code: def sample(
      logits, top_k=0, top_p=0.0, temperature=1.0, vocab_size=None
  ):
      """Sample and generate a token.
      Note: logits has the dimension [b, v] where b is the batch size
            and v is the vocabulary size.
      If vocab_size is provided, we will make sure the sample that is
      generated is in [0, vocab-size). This will avoid out of vocabulary
      generations due to padding.
      """
  
      # Check logits for consistency.
      assert (
          logits.ndim == 2
      ), "expected the logits to be of [b, v] shape."
      assert (
          logits.type() == "torch.cuda.FloatTensor"
      ), "input logits should be floats."
  
      # Greedy is just simple argmax.
      if top_k == 1:
          assert (
              top_p == 0.0
          ), "cannot set both greedy and top-p samplings."
          samples = torch.argmax(logits, dim=-1)
  
      # Top-k or top-p sampling.
      else:
          # Clone so we do not modify the inputs,
          logits = logits.clone()
          # Apply temperature in place.
          if temperature != 1.0:
              logits.div_(temperature)
  
          if top_k > 1:
              assert (
                  top_p == 0.0
              ), "cannot set both top-k and top-p samplings."
              assert top_k <= logits.size(
                  1
              ), "top-k is larger than logit size."
              if vocab_size:
                  assert (
                      top_k < vocab_size
                  ), "top-k is larger than vocab size."
              modify_logits_for_top_k_filtering(logits, top_k)
  
          elif top_p > 0.0:
              assert top_p <= 1.0, "top-p should be in (0, 1]."
              modify_logits_for_top_p_filtering(logits, top_p)
  
          # After filtering, we need to recalculate the distribution.
          probs = logits.softmax(dim=-1)
          samples = torch.multinomial(probs, num_samples=1).view(-1)
  
      # If vocab size is provided, make sure the samples are in
      # in the range [0, vocab-size).
      if vocab_size:
          samples = torch.clamp(samples, min=0, max=(vocab_size - 1))
  
      return samples
  
imports: import torch
label_desc: Use image node_id 3 for calling a global function with example usage: sample(logits, top_k, top_p, temperature, vocab_size) and returns: samples