code: def backward(ctx, grad_output):
      (indices,) = ctx.saved_tensors
      assert grad_output.ndim >= 2
      other_shape = grad_output.shape[1:]
      grad_output = rearrange(grad_output, "b ... -> b (...)")
      grad_input = torch.zeros(
          [ctx.first_axis_dim, grad_output.shape[1]],
          device=grad_output.device,
          dtype=grad_output.dtype,
      )
      # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
      # grad_input[indices] = grad_output
      grad_input.scatter_(
          0,
          repeat(indices, "z -> z d", d=grad_output.shape[1]),
          grad_output,
      )
      return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
  
file_path: flash-attention/flash_attn/bert_padding.py
imports: import torch
  import torch.nn.functional
  from einops import rearrange, repeat
label_desc: Use image node_id 2 for calling the IndexFirstAxis obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: grad_input,  None