class_docstr: An operator optimized for very small values of K (``K <= 32``)         and f32 pre-Ampere as it does not use TensorCores.
  Only supports contiguous inputs in BMK format, so an extra reshape         or contiguous call might be done.
  
  :Deprecated:
  
      This operator is deprecated and should not be used in new code
class_docstr_tok: ["An","operator","optimized","for","very","small","values","of","K","(","``","K","<","=","32","``",")","and","f32","pre-Ampere","as","it","does","not","use","TensorCores",".","Only","supports","contiguous","inputs","in","BMK","format",",","so","an","extra","reshape","or","contiguous","call","might","be","done",".",":","Deprecated",":","This","operator","is","deprecated","and","should","not","be","used","in","new","code"]
code: def apply(
      cls, inp: Inputs, needs_gradient: bool
  ) -> Tuple[torch.Tensor, Optional[Context]]:
      if inp.scale is not None:
          raise NotImplementedError("Unsupport custom scale")
      num_heads = inp.query.shape[2]
      query = _bmhk2bmk_contiguous(inp.query)
      key = _bmhk2bmk_contiguous(inp.key)
      value = _bmhk2bmk_contiguous(inp.value)
  
      out, lse, rng_seed, rng_offset = cls.OPERATOR(
          query=query,
          key=key,
          value=value,
          compute_logsumexp=needs_gradient,
          attn_bias=_get_tensor_bias_bmk(inp.attn_bias),
          p=inp.p,
      )
      out = bmk2bmhk(out, num_heads)
      lse = lse.reshape(
          [lse.shape[0] // num_heads, num_heads, lse.shape[1]]
      )
      if not needs_gradient:
          return out, None
      ctx = Context(out=out, lse=lse)
      if inp.p != 0.0:
          ctx.op_bw = BwOp
          ctx.rng_state = torch.tensor(
              [rng_seed, rng_offset], dtype=torch.int64, device="cpu"
          )
      return out, ctx
  
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 FwOp obj's underlying member method code with example usage: obj.apply(cls, inp, needs_gradient) and returns: out,  ctx, out,  None