class_docstr: Custom checkpointed function to get sum-per-token from an input and a weight.
class_docstr_tok: ["Custom","checkpointed","function","to","get","sum-per-token","from","an","input","and","a","weight","."]
docstr: Recompute the forward sum and backward grad.
  
  Accumulate the grad to the right split of the full grad.
docstr_tok: ["Recompute","the","forward","sum","and","backward","grad",".","Accumulate","the","grad","to","the","right","split","of","the","full","grad","."]
code: def backward(ctx: Any, *args: Any) -> Any:
      """Recompute the forward sum and backward grad.
  
      Accumulate the grad to the right split of the full grad.
      """
      if DEBUG and dist.is_initialized() and dist.get_rank() == 0:
          print("DEBUG sum bwd")
      assert len(args) == 1
      # Gradients should already exist due to TargetScoreFunction's backward.
      assert ctx.kernel_obj.proj_weight.grad is not None
  
      # Get saved i, w, and maxs.
      i, w, tgt, maxs = ctx.saved_tensors
      assert i.requires_grad
      assert w.requires_grad
      assert maxs.requires_grad
      i = i.detach().requires_grad_(True)
      w = w.detach().requires_grad_(True)
      maxs = maxs.detach().requires_grad_(True)
  
      # Forward + backward again.
      with torch.enable_grad():
          sums = GetSumFunction.get_sum(
              i,
              w,
              tgt,
              maxs,
              ctx.w_idx,
              ctx.kernel_obj.fp_sum,
              ctx.kernel_obj.margin,
              ctx.kernel_obj.scale,
          )
      torch.autograd.backward(sums, *args)
  
      # Accumulate the grads.
      assert w.grad is not None
      with torch.no_grad():
          grads = torch.split(
              ctx.kernel_obj.proj_weight.grad, ctx.w_split_size
          )
          grads[ctx.w_idx].add_(w.grad)
      return i.grad, None, None, maxs.grad, None, None, None, None
  
imports: from typing import Any, Optional, Tuple
  import torch
  from torch import nn
  import torch.distributed
  import torch.nn.functional
label_desc: Use image node_id 3 for calling the GetSumFunction obj's underlying member method code with example usage: obj.backward(ctx) and returns: i,  None,  None,  maxs,  None,  None,  None,  None