code: def sample(logits, top_k=1, top_p=0.0, temperature=1.0):
      """Sample from top-k logits.
      Arguments:
          logits: Tensor of shape (batch_size, vocab_size)
      """
      if top_k == 1:  # Short-circuit for greedy decoding
          return logits.argmax(dim=-1)
      else:
          if top_p > 0.0:
              assert top_p <= 1.0, "top-p should be in (0, 1]."
          if top_k > 0:
              top_k = min(top_k, logits.size(-1))  # Safety check
              logits_top, indices = torch.topk(logits, top_k, dim=-1)
              if temperature != 1.0:
                  logits_top /= temperature
              modify_logits_for_top_p_filtering(logits_top, top_p)
              return indices[
                  torch.arange(indices.shape[0], device=indices.device),
                  torch.multinomial(
                      torch.softmax(logits_top, dim=-1), num_samples=1
                  ).squeeze(dim=-1),
              ]
          else:
              # Clone so that when we modify for top_p we don't change the original logits
              logits_top = (
                  logits / temperature
                  if temperature != 1.0
                  else logits.clone()
              )
              modify_logits_for_top_p_filtering(logits_top, top_p)
              return torch.multinomial(
                  torch.softmax(logits_top, dim=-1), num_samples=1
              ).squeeze(dim=-1)
  
imports: import gc
  import time
  from collections import namedtuple
  from dataclasses import dataclass, field
  from functools import partial
  from typing import Callable, Optional, Sequence, Union
  import torch
  import torch.nn.functional
  from einops import rearrange, repeat
  from torch import Tensor
  from torch.profiler import ProfilerActivity, profile, record_function
  from transformers.generation import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput
label_desc: Use image node_id 3 for calling a global function with example usage: sample(logits, top_k, top_p, temperature) and returns: logits, indices, torch