docstr: checkpoint_lvl (increasing lvl means slower but more memory saving):
      0: no recomputation in the bwd
      1: recompute gelu_out in the bwd
      2: recompute gelu_in and gelu_out in the bwd
  heuristic:
      -1: don't fuse gemm + gelu (separate kernel)
      0..4: use this heuristic for the algo section in the fused gemm + gelu
docstr_tok: ["checkpoint_lvl","(","increasing","lvl","means","slower","but","more","memory","saving",")",":","0",":","no","recomputation","in","the","bwd","1",":","recompute","gelu_out","in","the","bwd","2",":","recompute","gelu_in","and","gelu_out","in","the","bwd","heuristic",":","-1",":","do","n't","fuse","gemm","+","gelu","(","separate","kernel",")","0","..","4",":","use","this","heuristic","for","the","algo","section","in","the","fused","gemm","+","gelu"]
code: def __init__(
      self,
      in_features,
      intermediate_features,
      out_features=None,
      bias=True,
      checkpoint_lvl=0,
      bf16=False,
      heuristic=0,
      device=None,
      dtype=None,
  ):
      """
      checkpoint_lvl (increasing lvl means slower but more memory saving):
          0: no recomputation in the bwd
          1: recompute gelu_out in the bwd
          2: recompute gelu_in and gelu_out in the bwd
      heuristic:
          -1: don't fuse gemm + gelu (separate kernel)
          0..4: use this heuristic for the algo section in the fused gemm + gelu
      """
      assert checkpoint_lvl in [0, 1, 2]
      factory_kwargs = {"device": device, "dtype": dtype}
      super().__init__()
      if out_features is None:
          out_features = in_features
      assert (
          bias == True
      ), "DenseGeluDense module without bias is currently not supported"
      self.checkpoint_lvl = checkpoint_lvl
      self.bf16 = bf16
      self.heuristic = heuristic
      self.fc1 = nn.Linear(
          in_features,
          intermediate_features,
          bias=bias,
          **factory_kwargs,
      )
      self.fc2 = nn.Linear(
          intermediate_features,
          out_features,
          bias=bias,
          **factory_kwargs,
      )
  
imports: import numpy
  import torch
  import torch.nn
  import torch.nn.functional
  from torch.cuda.amp import custom_bwd, custom_fwd
  import fused_dense_lib
  from src.ops.gelu_activation import gelu_bwd
label_desc: Use image node_id 1 to create a new FusedDenseGeluDenseTD object from inherited base classes: nn with example: obj = FusedDenseGeluDenseTD(in_features, intermediate_features, out_features, bias, checkpoint_lvl, bf16, heuristic, device, dtype)