code: def apply(
      cls, ctx: Context, inp: Inputs, grad: torch.Tensor
  ) -> Gradients:
      num_heads = grad.shape[2]
      grad = _bmhk2bmk_contiguous(grad)
      query = _bmhk2bmk_contiguous(inp.query)
      key = _bmhk2bmk_contiguous(inp.key)
      value = _bmhk2bmk_contiguous(inp.value)
      out = _bmhk2bmk_contiguous(ctx.out)
  
      rng_seed = rng_offset = 0
      if inp.p != 0.0:
          if (
              ctx.rng_state is None
              or ctx.rng_state.dtype != torch.int64
              or ctx.rng_state.device.type != "cpu"
              or ctx.rng_state.shape != (2,)
          ):
              raise NotImplementedError(
                  f"Invalid rng_state: {ctx.rng_state}"
              )
          rng_seed, rng_offset = ctx.rng_state.tolist()
      grad_q, grad_k, grad_v = cls.OPERATOR(
          grad,
          query,
          key,
          value,
          # LSE: BHM -> (BH)M
          ctx.lse.flatten(end_dim=-2),
          out,
          _get_tensor_bias_bmk(inp.attn_bias),
          inp.p,
          rng_seed,
          rng_offset,
      )
      return Gradients(
          dq=bmk2bmhk(grad_q, num_heads),
          dk=bmk2bmhk(grad_k, num_heads),
          dv=bmk2bmhk(grad_v, num_heads),
      )
  
file_path: xformers/xformers/ops/fmha/small_k.py
imports: from typing import Any, List, Mapping, Optional, Set, Tuple, Union
  import torch
  from ..common import get_xformers_operator, register_operator
  from .attn_bias import AttentionBias
  from .common import AttentionBwOpBase, AttentionFwOpBase, Context, Gradients, Inputs, bmk2bmhk
label_desc: Use image node_id 2 for calling the BwOp obj's underlying member method code with example usage: obj.apply(cls, ctx, inp, grad) and returns: Gradients