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: Forward function that computes the max, without saving activations.
docstr_tok: ["Forward","function","that","computes","the","max",",","without","saving","activations","."]
code: def forward(  # type: ignore
      ctx: Any,
      i: torch.Tensor,
      w: torch.Tensor,
      tgt: torch.Tensor,
      kernel_obj: "MemoryEfficientVocabOutput",
      w_idx: int,
      w_split_size: int,
      split_dim: int,
  ) -> torch.Tensor:
      """Forward function that computes the max, without saving activations."""
      if DEBUG and dist.is_initialized() and dist.get_rank() == 0:
          print("DEBUG max fwd")
      ctx.save_for_backward(i, w, tgt)
      ctx.kernel_obj = kernel_obj
      ctx.w_idx = w_idx
      ctx.w_split_size = w_split_size
      ctx.args = {}
      assert split_dim == 0
      # During forward, we use ``no_grad'' to avoid saving the activations.
      # The activations will be recomputed in backward below and freed
      # immediately after use. This saves the overall GPU peak memory of this layer.
      with torch.no_grad():
          return GetMaxFunction.get_max(
              i,
              w,
              tgt,
              w_idx,
              kernel_obj.fp_max,
              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 GetMaxFunction obj's underlying member method code with example usage: obj.forward(ctx, i, w, tgt, kernel_obj, w_idx, w_split_size, split_dim) and returns: GetMaxFunction