code: def GumbelCRFSemiring(temp):
      class ST(torch.autograd.Function):
          @staticmethod
          def forward(ctx, logits, dim):
              out = torch.nn.functional.one_hot(logits.max(-1)[1], dim)
              out = out.type_as(logits)
              ctx.save_for_backward(logits, out)
              return out
  
          @staticmethod
          def backward(ctx, grad_output):
              logits, out = ctx.saved_tensors
              with torch.enable_grad():
                  ret = torch.autograd.grad(
                      logits.softmax(-1), logits, out * grad_output
                  )[0]
              return ret, None
  
      class _GumbelCRFLogSumExp(torch.autograd.Function):
          @staticmethod
          def forward(ctx, input, dim):
              ctx.save_for_backward(input, torch.tensor(dim))
              return torch.logsumexp(input, dim=dim)
  
          @staticmethod
          def backward(ctx, grad_output):
              logits, dim = ctx.saved_tensors
              grad_input = None
              if ctx.needs_input_grad[0]:
  
                  def sample(ls):
                      update = (
                          ls
                          + torch.distributions.Gumbel(0, 1).sample(
                              (ls.shape[-1],)
                          )
                      ) / temp
                      out = ST.apply(update, ls.shape[-1])
                      return out
  
                  if dim == -1:
                      s = sample(logits)
                  else:
                      dim = dim if dim >= 0 else logits.dim() + dim
                      perm = [
                          i for i in range(logits.dim()) if i != dim
                      ] + [dim]
                      rev_perm = [
                          a
                          for a, b in sorted(
                              enumerate(perm), key=lambda a: a[1]
                          )
                      ]
                      s = sample(logits.permute(perm)).permute(rev_perm)
  
                  grad_input = grad_output.unsqueeze(dim).mul(s)
              return grad_input, None
  
      class _GumbelCRFSemiring(_BaseLog):
          @staticmethod
          def sum(xs, dim=-1):
              return _GumbelCRFLogSumExp.apply(xs, dim)
  
      return _GumbelCRFSemiring
  
file_path: pytorch-struct/torch_struct/semirings/sample.py
imports: import torch
  import torch.distributions
  from .semirings import _BaseLog
label_desc: Use image node_id 2 for calling a global function with example usage: GumbelCRFSemiring(temp) and returns: _GumbelCRFSemiring, out, ret,  None, torch, grad_input,  None, _GumbelCRFLogSumExp, out