code: def backward(ctx, grad_output):
      grad_output = grad_output.contiguous()
      checkpoint_lvl = ctx.checkpoint_lvl
      x, weight1, bias1, weight2, *rest = ctx.saved_tensors
      batch_shape, n = x.shape[:-1], x.shape[-1]
      batch_dim = np.prod(batch_shape)
      if checkpoint_lvl == 0:
          gelu_in, output1 = rest
      elif checkpoint_lvl == 1:
          (gelu_in,) = rest
          output1 = F.gelu(gelu_in, approximate="tanh")
      elif checkpoint_lvl == 2:
          # bias1, = rest
          if ctx.heuristic == -1:
              gelu_in = fused_dense_cuda.linear_bias_forward(
                  x.reshape(batch_dim, n), weight1, bias1
              )
              output1 = F.gelu(gelu_in, approximate="tanh")
          else:
              output1, gelu_in = fused_dense_cuda.linear_gelu_forward(
                  x.reshape(batch_dim, n),
                  weight1,
                  bias1,
                  True,
                  ctx.heuristic,
              )
  
      if ctx.heuristic == -1:
          grad_output = grad_output.reshape(
              batch_dim, grad_output.shape[-1]
          )
          # grad_output1, grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_backward(output1, weight2, grad_output)
          grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(
              output1, grad_output
          )
          # grad_gelu = matmul_dgelu(grad_output, weight2, gelu_in)
          grad_output1 = grad_output @ weight2
          with torch.jit.fuser("fuser2"):
              grad_gelu = gelu_bwd(grad_output1, gelu_in)
          (
              grad_input,
              grad_weight1,
              grad_bias1,
          ) = fused_dense_cuda.linear_bias_backward(
              x.reshape(batch_dim, n), weight1, grad_gelu
          )
          # with torch.jit.fuser('fuser2'):
          #     grad_gelu, grad_bias1 = bias_gelu_back(grad_output1, gelu_in, bias1)
          # grad_input = grad_gelu @ weight1
          # grad_weight1 = grad_gelu.reshape(batch_dim, -1).T @ x.reshape(batch_dim, n)
          # grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward(
          #     x.reshape(batch_dim, n), weight1, grad_gelu
          # )
      else:
          (
              grad_input,
              grad_weight1,
              grad_bias1,
              grad_weight2,
              grad_bias2,
          ) = fused_dense_cuda.linear_gelu_linear_backward(
              x.reshape(batch_dim, n),
              gelu_in,
              output1,
              weight1,
              weight2,
              grad_output.reshape(batch_dim, grad_output.shape[-1]),
              ctx.heuristic,
          )
      # grad_output = grad_output.reshape(batch_dim, grad_output.shape[-1])
      # # grad_output1, grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_backward(output1, weight2, grad_output)
      # grad_weight2, grad_bias2 = fused_dense_cuda.linear_bias_wgrad(output1, grad_output)
      # grad_gelu = matmul_dgelu(grad_output, weight2, gelu_in)
      # grad_input, grad_weight1, grad_bias1 = fused_dense_cuda.linear_bias_backward(
      #     x.reshape(batch_dim, n), weight1, grad_gelu
      # )
      return (
          grad_input.reshape_as(x),
          grad_weight1,
          grad_bias1,
          grad_weight2,
          grad_bias2,
          None,
          None,
          None,
      )
  
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 2 for calling the FusedDenseGeluDenseFuncTD obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_input,  grad_weight1,  grad_bias1,  grad_weight2,  grad_bias2,  None,  None,  None