code: def sample_from_logits(a):
      # returns a sampled element and the associated probability
      # since cross entropy is run during training we expect logits
      # to be probabilities yet.
      probs = torch.cumsum(F.softmax(a, dim=-1), dim=-1)
      random_vals = (
          torch.rand(probs.shape[0]).unsqueeze(-1).to(a.device)
      )
      n_classes = a.shape[-1]
      new_a_idx = torch.argmax(1.0 * (probs > random_vals), dim=-1)
      index_bias = (
          torch.arange(0, len(new_a_idx)).to(a.device) * n_classes
      )
      probs = torch.take(probs, new_a_idx + index_bias)
      # new_a = F.one_hot(new_a_idx, n_classes)
      return new_a_idx, probs
  
imports: import math
  import torch
  import torch.nn.functional
  from open_alpha_tensor.core.modules.attention import AlphaMultiHeadAttention
label_desc: Use image node_id 1 for calling a global function with example usage: sample_from_logits(a) and returns: new_a_idx,  probs