class_docstr: MultiScaleAttention
class_docstr_tok: ["MultiScaleAttention"]
code: def __init__(
      self,
      dim,
      num_heads=8,
      qkv_bias=False,
      drop_rate=0.0,
      kernel_q=(1, 1, 1),
      kernel_kv=(1, 1, 1),
      stride_q=(1, 1, 1),
      stride_kv=(1, 1, 1),
      norm_layer=LayerNorm,
      has_cls_embed=True,
      # Options include `conv`, `avg`, and `max`.
      mode="conv",
  ):
      super().__init__()
      self.drop_rate = drop_rate
      self.num_heads = num_heads
      head_dim = dim // num_heads
      self.scale = head_dim**-0.5
      self.has_cls_embed = has_cls_embed
      padding_q = [int(q // 2) for q in kernel_q]
      padding_kv = [int(kv // 2) for kv in kernel_kv]
  
      self.qkv = Linear(
          dim, dim * 3, bias=qkv_bias
      )  # pylint: disable=too-many-function-args,unexpected-keyword-arg
      self.proj = Linear(
          dim, dim
      )  # pylint: disable=too-many-function-args
      # self.show_attn = nn.Softmax(dim=-1)
      if drop_rate > 0.0:
          self.proj_drop = nn.Dropout(drop_rate)
  
      # A = Q*K^T
      self.matmul1 = Einsum("bhid,bhjd->bhij")
      # attn = A*V
      self.matmul2 = Einsum("bhij,bhjd->bhid")
  
      self.softmax = Softmax(
          dim=-1
      )  # pylint: disable=unexpected-keyword-arg
  
      # Skip pooling with kernel and stride size of (1, 1, 1).
      if numpy.prod(kernel_q) == 1 and numpy.prod(stride_q) == 1:
          kernel_q = ()
      if numpy.prod(kernel_kv) == 1 and numpy.prod(stride_kv) == 1:
          kernel_kv = ()
      pool_q = pool_k = pool_v = norm_q = norm_k = norm_v = None
      if mode == "avg":
          pool_q = (
              nn.AvgPool3d(
                  kernel_q, stride_q, padding_q, ceil_mode=False
              )
              if len(kernel_q) > 0
              else None
          )
          pool_k = (
              nn.AvgPool3d(
                  kernel_kv, stride_kv, padding_kv, ceil_mode=False
              )
              if len(kernel_kv) > 0
              else None
          )
          pool_v = (
              nn.AvgPool3d(
                  kernel_kv, stride_kv, padding_kv, ceil_mode=False
              )
              if len(kernel_kv) > 0
              else None
          )
      elif mode == "max":
          pool_q = (
              nn.MaxPool3d(
                  kernel_q, stride_q, padding_q, ceil_mode=False
              )
              if len(kernel_q) > 0
              else None
          )
          pool_k = (
              nn.MaxPool3d(
                  kernel_kv, stride_kv, padding_kv, ceil_mode=False
              )
              if len(kernel_kv) > 0
              else None
          )
          pool_v = (
              nn.MaxPool3d(
                  kernel_kv, stride_kv, padding_kv, ceil_mode=False
              )
              if len(kernel_kv) > 0
              else None
          )
      elif mode == "conv":
          pool_q = (
              Conv3d(  # pylint: disable=too-many-function-args,unexpected-keyword-arg
                  head_dim,
                  head_dim,
                  kernel_q,
                  stride=stride_q,
                  padding=padding_q,
                  groups=head_dim,
                  bias=False,
              )
              if len(kernel_q) > 0
              else None
          )
          norm_q = norm_layer(head_dim) if len(kernel_q) > 0 else None
          pool_k = (
              Conv3d(  # pylint: disable=too-many-function-args,disable=unexpected-keyword-arg
                  head_dim,
                  head_dim,
                  kernel_kv,
                  stride=stride_kv,
                  padding=padding_kv,
                  groups=head_dim,
                  bias=False,
              )
              if len(kernel_kv) > 0
              else None
          )
          norm_k = norm_layer(head_dim) if len(kernel_kv) > 0 else None
          pool_v = (
              Conv3d(  # pylint: disable=too-many-function-args,unexpected-keyword-arg
                  head_dim,
                  head_dim,
                  kernel_kv,
                  stride=stride_kv,
                  padding=padding_kv,
                  groups=head_dim,
                  bias=False,
              )
              if len(kernel_kv) > 0
              else None
          )
          norm_v = norm_layer(head_dim) if len(kernel_kv) > 0 else None
      else:
          raise NotImplementedError(f"Unsupported model {mode}")
  
      self.atn_pool_q = AttentionPool(
          pool=pool_q,
          # thw_shape=thw_shape,
          has_cls_embed=self.has_cls_embed,
          norm=norm_q,
      )
      self.atn_pool_k = AttentionPool(
          pool=pool_k,
          # thw_shape=thw_shape,
          has_cls_embed=self.has_cls_embed,
          norm=norm_k,
      )
      self.atn_pool_v = AttentionPool(
          pool=pool_v,
          # thw_shape=thw_shape,
          has_cls_embed=self.has_cls_embed,
          norm=norm_v,
      )
      self.attn_cam = None
      self.attn = None
      self.v = None
      self.v_cam = None
      self.attn_gradients = None
  
imports: import math
  from functools import partial
  from typing import Tuple, List
  import numpy
  import torch
  from torch import nn
  from torch.nn.init import trunc_normal_
  from einops import rearrange
  from towhee.models.layers.layers_with_relprop import GELU, Linear, Dropout, Conv2d, Conv3d, Softmax, Sigmoid, LayerNorm, Einsum, Add, Clone, MaxPool3d, IndexSelect
label_desc: Use image node_id 1 to create a new MultiScaleAttention object from inherited base classes: nn with example: obj = MultiScaleAttention(dim, num_heads, qkv_bias, drop_rate, kernel_q, kernel_kv, stride_q, stride_kv, norm_layer, has_cls_embed, mode)