class_docstr: Custom checkpointed function to get max-per-token from an input and a weight
class_docstr_tok: ["Custom","checkpointed","function","to","get","max-per-token","from","an","input","and","a","weight"]
docstr: Throughout this code:
  
    i: input data with shape = (split-of-tokens, d_model)
    w: weight data with shape = (split-of-vocabs, d_model)
    tgt: target prediction data with shape = (split-of-tokens,)
docstr_tok: ["Throughout","this","code",":","i",":","input","data","with","shape","=","(","split-of-tokens",",","d_model",")","w",":","weight","data","with","shape","=","(","split-of-vocabs",",","d_model",")","tgt",":","target","prediction","data","with","shape","=","(","split-of-tokens",",",")"]
code: def get_max(
      i: torch.Tensor,
      w: torch.Tensor,
      tgt: torch.Tensor,
      w_idx: int,
      full_precision: bool,
      margin: float,
      scale: Optional[float],
  ) -> torch.Tensor:
      """
      Throughout this code:
  
        i: input data with shape = (split-of-tokens, d_model)
        w: weight data with shape = (split-of-vocabs, d_model)
        tgt: target prediction data with shape = (split-of-tokens,)
      """
      if scale is not None:
          _m = lmcl_matmul(i, w, tgt, w_idx, margin, scale)
      else:
          _m = torch.matmul(i, w.T)
      if full_precision:
          _m = _m.float()
      _m = _m.max(dim=1)[0]
      return _m
  
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 GetMaxFunction obj's underlying member method code with example usage: obj.get_max(i, w, tgt, w_idx, full_precision, margin, scale) and returns: _m