class_docstr: A module which performs QKV attention and splits in a different order.
code: def __init__(self, n_heads):
      super().__init__()
      self.n_heads = n_heads
  
imports: import math
  import numpy
  import torch
  import torch.nn
  from torch.cuda.amp import custom_bwd, custom_fwd
label_desc: Use image node_id 1 to create a new QKVStableAttention object from inherited base classes: nn with example: obj = QKVStableAttention(n_heads)