code: def argmax() -> Sampler:
      def sample(logits: Logits) -> Sample:
          return Sample(
              logits=logits, tokens=torch.argmax(logits, dim=-1)
          )
  
      return sample
  
imports: from dataclasses import dataclass
  from typing import Callable
  import torch
label_desc: Use image node_id 2 for calling a global function with example usage: argmax() and returns: sample, Sample