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: Forward function that computes the sum, without saving activations.
docstr_tok: ["Forward","function","that","computes","the","sum",",","without","saving","activations","."]
code: def forward(  # type: ignore
      ctx: Any,
      i: torch.Tensor,
      w: torch.Tensor,
      tgt: torch.Tensor,
      maxs: torch.Tensor,
      kernel_obj: "MemoryEfficientVocabOutput",
      w_idx: int,
      w_split_size: int,
      split_dim: int,
  ) -> torch.Tensor:
      """Forward function that computes the sum, without saving activations."""
      if DEBUG and dist.is_initialized() and dist.get_rank() == 0:
          print("DEBUG sum fwd")
      ctx.save_for_backward(i, w, tgt, maxs)
      ctx.kernel_obj = kernel_obj
      ctx.w_idx = w_idx
      ctx.w_split_size = w_split_size
      assert split_dim == 0
      with torch.no_grad():
          return GetSumFunction.get_sum(
              i,
              w,
              tgt,
              maxs,
              w_idx,
              kernel_obj.fp_sum,
              kernel_obj.margin,
              kernel_obj.scale,
          )
  
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 2 for calling the GetSumFunction obj's underlying member method code with example usage: obj.forward(ctx, i, w, tgt, maxs, kernel_obj, w_idx, w_split_size, split_dim) and returns: GetSumFunction