code: def backward(ctx, grad_output):
      logits, dim = ctx.saved_tensors
      grad_input = None
      if ctx.needs_input_grad[0]:
  
          def sample(ls):
              pre_shape = ls.shape
              draws = torch.multinomial(
                  ls.softmax(-1).view(-1, pre_shape[-1]), 1, True
              )
              draws.squeeze(1)
              return (
                  torch.nn.functional.one_hot(draws, pre_shape[-1])
                  .view(*pre_shape)
                  .type_as(ls)
              )
  
          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
  
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 the _SampledLogSumExp obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_input,  None, torch