docstr: checkpoint_lvl:
  0: no recomputation in the bwd
  1: recompute gelu_out in the bwd
  2: recompute gelu_in and gelu_out in the bwd
docstr_tok: ["checkpoint_lvl",":","0",":","no","recomputation","in","the","bwd","1",":","recompute","gelu_out","in","the","bwd","2",":","recompute","gelu_in","and","gelu_out","in","the","bwd"]
code: def forward(
      ctx,
      x,
      weight1,
      bias1,
      weight2,
      bias2,
      checkpoint_lvl=0,
      heuristic=0,
  ):
      """checkpoint_lvl:
      0: no recomputation in the bwd
      1: recompute gelu_out in the bwd
      2: recompute gelu_in and gelu_out in the bwd
      """
      assert checkpoint_lvl in [0, 1, 2]
      assert -1 <= heuristic <= 4
      x = x.contiguous()
      weight1 = weight1.contiguous()
      bias1 = bias1.contiguous()
      weight2 = weight2.contiguous()
      bias2 = bias2.contiguous()
      batch_shape, n = x.shape[:-1], x.shape[-1]
      batch_dim = np.prod(batch_shape)
      assert (
          batch_dim <= 64 * 1024
      ), "fused_dense only supports dimension at most 64k"
      # output1, output2, gelu_in = fused_dense_cuda.linear_gelu_linear_forward(
      #     x.reshape(batch_dim, n), weight1, bias1, weight2, bias2
      # )
      # gelu_in = fused_dense_cuda.linear_bias_forward(x.reshape(batch_dim, n), weight1, bias1)
      # output1 = F.gelu(gelu_in, approximate='tanh')
      save_gelu_in = checkpoint_lvl != 2
      output1, *rest = fused_dense_cuda.linear_gelu_forward(
          x.reshape(batch_dim, n),
          weight1,
          bias1,
          save_gelu_in,
          heuristic,
      )
      if save_gelu_in:
          gelu_in = rest[0]
      output2 = fused_dense_cuda.linear_bias_forward(
          output1, weight2, bias2
      )
      ctx.checkpoint_lvl = checkpoint_lvl
      ctx.heuristic = heuristic
      if checkpoint_lvl == 0:
          ctx.save_for_backward(x, weight1, weight2, gelu_in, output1)
      elif checkpoint_lvl == 1:
          ctx.save_for_backward(x, weight1, weight2, gelu_in)
      elif checkpoint_lvl == 2:
          ctx.save_for_backward(x, weight1, weight2, bias1)
      return output2.reshape(*batch_shape, output2.shape[-1]), x
  
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 for calling the FusedDenseResGeluDenseFunc obj's underlying member method code with example usage: obj.forward(ctx, x, weight1, bias1, weight2, bias2, checkpoint_lvl, heuristic) and returns: output2,  x