code: def __init__(
      self,
      input_dim: int,
      num_heads: int = 1,
      *,
      bias: bool = True,
      dropout: float = 0.0,
      qk_scale: Optional[float] = None,
      kv_same: Optional[bool] = None,
      qkv_bias_same: bool = True,
      is_self_attention: bool = False,
      k_dim: Optional[int] = None,
      v_dim: Optional[int] = None,
      embed_dim: Optional[int] = None,
      activation: Optional[str] = None,
      activation_config: Optional[Dict[str, Any]] = None,
      out_linear_config: Optional[Dict[str, Any]] = None,
      reduction_ratio: Optional[int] = None,
      hook: Optional[IAttentionHook] = None,
  ):
      super().__init__()
      self.input_dim = input_dim
      if kv_same is None:
          if k_dim is not None and v_dim is not None and k_dim != v_dim:
              kv_same = False
          else:
              kv_same = True
      self.kv_same = kv_same
      has_reduction = (
          reduction_ratio is not None and reduction_ratio > 1
      )
      self.qkv_same = is_self_attention and not has_reduction
      if not is_self_attention:
          self.k_dim = k_dim or input_dim
          self.v_dim = v_dim or self.k_dim
      else:
          if k_dim is not None and k_dim != input_dim:
              raise ValueError(
                  "self attention is used but `k_dim` != `input_dim`"
              )
          if v_dim is not None and v_dim != input_dim:
              raise ValueError(
                  "self attention is used but `v_dim` != `input_dim`"
              )
          self.k_dim = self.v_dim = input_dim
      self.embed_dim = embed_dim or input_dim
  
      self.num_heads = num_heads
      self.head_dim = self.embed_dim // num_heads
      self.scaling = qk_scale or float(self.head_dim) ** 0.5
      if self.head_dim * num_heads != self.embed_dim:
          raise ValueError(
              "`embed_dim` must be divisible by `num_heads`"
          )
  
      if self.qkv_same:
          self.kv_w = self.q_w = self.k_w = self.v_w = None
          self.in_w = nn.Parameter(
              torch.empty(3 * self.embed_dim, input_dim)
          )
          nn.init.trunc_normal_(self.in_w, std=0.02)
      elif kv_same:
          self.in_w = self.k_w = self.v_w = None
          self.q_w = nn.Parameter(
              torch.empty(self.embed_dim, input_dim)
          )
          self.kv_w = nn.Parameter(
              torch.empty(2 * self.embed_dim, input_dim)
          )
          nn.init.trunc_normal_(self.q_w, std=0.02)
          nn.init.trunc_normal_(self.kv_w, std=0.02)
      else:
          self.in_w = None
          self.q_w = nn.Parameter(
              torch.empty(self.embed_dim, input_dim)
          )
          self.k_w = nn.Parameter(
              torch.empty(self.embed_dim, self.k_dim)
          )
          self.v_w = nn.Parameter(
              torch.empty(self.embed_dim, self.v_dim)
          )
          nn.init.xavier_uniform_(self.q_w)
          nn.init.xavier_uniform_(self.k_w)
          nn.init.xavier_uniform_(self.v_w)
      if not bias:
          self.q_bias = self.k_bias = self.v_bias = None
          self.kv_bias = self.qkv_bias = None
      elif not qkv_bias_same:
          self.kv_bias = self.qkv_bias = None
          self.q_bias = nn.Parameter(torch.zeros(self.embed_dim))
          self.k_bias = nn.Parameter(torch.zeros(self.embed_dim))
          self.v_bias = nn.Parameter(torch.zeros(self.embed_dim))
      elif self.qkv_same or not kv_same:
          self.q_bias = self.k_bias = self.v_bias = self.kv_bias = None
          self.qkv_bias = nn.Parameter(torch.zeros(3 * self.embed_dim))
      else:
          self.k_bias = self.v_bias = self.qkv_bias = None
          self.q_bias = nn.Parameter(torch.zeros(self.embed_dim))
          self.kv_bias = nn.Parameter(torch.zeros(2 * self.embed_dim))
  
      self.out_linear = HijackCustomLinear(
          self.embed_dim,
          input_dim,
          **(out_linear_config or {}),
      )
  
      self.dropout = dropout
      self.activation = Activation.make(activation, activation_config)
  
      if not has_reduction:
          self.reduction = None
      else:
          self.reduction = nn.Sequential(
              Conv2d(
                  self.embed_dim,
                  self.embed_dim,
                  kernel_size=reduction_ratio,  # type: ignore
                  stride=reduction_ratio,  # type: ignore
                  padding=0,
              ),
              Lambda(lambda t: t.flatten(2).transpose(1, 2)),
              nn.LayerNorm(self.embed_dim),
          )
  
      # hook
      self.hook = hook
  
imports: import math
  import torch
  import numpy
  import torch.nn
  import torch.nn.functional
  from torch import Tensor
  from typing import Any
  from typing import Dict
  from typing import Type
  from typing import Tuple
  from typing import Optional
  from typing import NamedTuple
  from functools import partial
  from torch.nn import Module
  from cftool.misc import safe_execute
  from cftool.misc import WithRegister
  from .convs import conv_nd
  from .convs import Conv2d
  from .utils import zero_module
  from .common import Lambda
  from .hooks import IAttentionHook
  from .hijacks import IAttention
  from .hijacks import IHijackMixin
  from .hijacks import HijackConv2d
  from .hijacks import HijackLinear
  from .hijacks import HijackCustomLinear
  from .activations import Activation
  from ...misc.toolkit import sdp_attn
  from ...misc.toolkit import gradient_checkpoint
label_desc: Use image node_id 1 to create a new Attention object from inherited base classes: Module, IAttention, IHijackMixin, WithRegister with example: obj = Attention(input_dim, num_heads)