code: def backward(ctx, grad_output):
      # Retreive tensors from the forward path.
      (
          softmax,
          target_mask,
          masked_target_1d,
          label_smoothing,
          vocab_size,
      ) = ctx.saved_tensors
  
      label_smoothing = label_smoothing.item()
      vocab_size = vocab_size.item()
  
      # All the inputs have softmax as thier gradient.
      grad_input = softmax
      # For simplicity, work with the 2D gradient.
      partition_vocab_size = softmax.size()[-1]
      grad_2d = grad_input.view(-1, partition_vocab_size)
  
      # Add the gradient from matching classes.
      arange_1d = torch.arange(
          start=0, end=grad_2d.size()[0], device=grad_2d.device
      )
  
      softmax_update = 1.0 - target_mask.view(-1).float()
  
      if label_smoothing > 0:
          smoothing = label_smoothing * vocab_size / (vocab_size - 1)
          grad_2d[arange_1d, masked_target_1d] -= (
              1.0 - smoothing
          ) * softmax_update
          average_grad = 1 / vocab_size
          grad_2d[arange_1d, :] -= smoothing * average_grad
      else:
          grad_2d[arange_1d, masked_target_1d] -= softmax_update
  
      # Finally elementwise multiplication with the output gradients.
      grad_input.mul_(grad_output.unsqueeze(dim=-1))
  
      return grad_input, None, None
  
imports: import torch
label_desc: Use image node_id 2 for calling the _VocabParallelCrossEntropy obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_input,  None,  None