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","."]
code: def get_sum(
      i: torch.Tensor,
      w: torch.Tensor,
      tgt: torch.Tensor,
      maxs: torch.Tensor,
      w_idx: int,
      full_precision: bool,
      margin: float,
      scale: Optional[float],
  ) -> torch.Tensor:
      if scale is not None:
          _s = lmcl_matmul(i, w, tgt, w_idx, margin, scale)
      else:
          _s = torch.matmul(i, w.T)
      if full_precision:
          _s = _s.float()
      _s = (_s - maxs.reshape(-1, 1)).exp().sum(dim=1)
      return _s
  
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 1 for calling the GetSumFunction obj's underlying member method code with example usage: obj.get_sum(i, w, tgt, maxs, w_idx, full_precision, margin, scale) and returns: _s