class_docstr: Implementation of a multiscale attention block. Compare to a conventional attention
  block, a multiscale attention block optionally supports pooling (either
  before or after qkv projection). If pooling is not used, a multiscale attention
  block is equivalent to a conventional attention block.
  
  ::
                                 Input
                                   |
                  |----------------|-----------------|
                  ↓                ↓                 ↓
                Linear           Linear            Linear
                  &                &                 &
               Pool (Q)         Pool (K)          Pool (V)
                  → -------------- ←                 |
                           ↓                         |
                     MatMul & Scale                  |
                           ↓                         |
                        Softmax                      |
                           → ----------------------- ←
                                       ↓
                                 MatMul & Scale
                                       ↓
                                    DropOut
docstr: Args:
      dim (int): Input feature dimension.
      num_heads (int): Number of heads in the attention layer.
      qkv_bias (bool): If set to False, the qkv layer will not learn an additive
          bias. Default: False.
      dropout_rate (float): Dropout rate.
      kernel_q (_size_3_t): Pooling kernel size for q. If both pooling kernel
          size and pooling stride size are 1 for all the dimensions, pooling is
          disabled.
      kernel_kv (_size_3_t): Pooling kernel size for kv. If both pooling kernel
          size and pooling stride size are 1 for all the dimensions, pooling is
          disabled.
      stride_q (_size_3_t): Pooling kernel stride for q.
      stride_kv (_size_3_t): Pooling kernel stride for kv.
      norm_layer (nn.Module): Normalization layer used after pooling.
      has_cls_embed (bool): If set to True, the first token of the input tensor
          should be a cls token. Otherwise, the input tensor does not contain a
          cls token. Pooling is not applied to the cls token.
      pool_mode (str): Pooling mode. Option includes "conv" (learned pooling), "avg"
          (average pooling), and "max" (max pooling).
      pool_first (bool): If set to True, pool is applied before qkv projection.
          Otherwise, pool is applied after qkv projection. Default: False.
      residual_pool (bool): If set to True, use Improved Multiscale Vision
          Transformer's pooling residual connection.
      depthwise_conv (bool): Whether use depthwise or full convolution for pooling.
      bias_on (bool): Whether use biases for linear layers.
      separate_qkv (bool): Whether to use separate or one layer for qkv projections.
docstr_tok: ["Args",":","dim","(","int",")",":","Input","feature","dimension",".","num_heads","(","int",")",":","Number","of","heads","in","the","attention","layer",".","qkv_bias","(","bool",")",":","If","set","to","False",",","the","qkv","layer","will","not","learn","an","additive","bias",".","Default",":","False",".","dropout_rate","(","float",")",":","Dropout","rate",".","kernel_q","(","_size_3_t",")",":","Pooling","kernel","size","for","q",".","If","both","pooling","kernel","size","and","pooling","stride","size","are","1","for","all","the","dimensions",",","pooling","is","disabled",".","kernel_kv","(","_size_3_t",")",":","Pooling","kernel","size","for","kv",".","If","both","pooling","kernel","size","and","pooling","stride","size","are","1","for","all","the","dimensions",",","pooling","is","disabled",".","stride_q","(","_size_3_t",")",":","Pooling","kernel","stride","for","q.","stride_kv","(","_size_3_t",")",":","Pooling","kernel","stride","for","kv",".","norm_layer","(","nn.Module",")",":","Normalization","layer","used","after","pooling",".","has_cls_embed","(","bool",")",":","If","set","to","True",",","the","first","token","of","the","input","tensor","should","be","a","cls","token",".","Otherwise",",","the","input","tensor","does","not","contain","a","cls","token",".","Pooling","is","not","applied","to","the","cls","token",".","pool_mode","(","str",")",":","Pooling","mode",".","Option","includes","``","conv","''","(","learned","pooling",")",",","``","avg","''","(","average","pooling",")",",","and","``","max","''","(","max","pooling",")",".","pool_first","(","bool",")",":","If","set","to","True",",","pool","is","applied","before","qkv","projection",".","Otherwise",",","pool","is","applied","after","qkv","projection",".","Default",":","False",".","residual_pool","(","bool",")",":","If","set","to","True",",","use","Improved","Multiscale","Vision","Transformer","'s","pooling","residual","connection",".","depthwise_conv","(","bool",")",":","Whether","use","depthwise","or","full","convolution","for","pooling",".","bias_on","(","bool",")",":","Whether","use","biases","for","linear","layers",".","separate_qkv","(","bool",")",":","Whether","to","use","separate","or","one","layer","for","qkv","projections","."]
code: def __init__(
      self,
      dim: int,
      num_heads: int = 8,
      qkv_bias: bool = False,
      dropout_rate: float = 0.0,
      kernel_q: _size_3_t = (1, 1, 1),
      kernel_kv: _size_3_t = (1, 1, 1),
      stride_q: _size_3_t = (1, 1, 1),
      stride_kv: _size_3_t = (1, 1, 1),
      norm_layer=nn.LayerNorm,
      has_cls_embed: bool = True,
      pool_mode: str = "conv",
      pool_first: bool = False,
      residual_pool: bool = True,
      depthwise_conv: bool = True,
      bias_on: bool = True,
      separate_qkv: bool = True,
  ) -> None:
      """
      Args:
          dim (int): Input feature dimension.
          num_heads (int): Number of heads in the attention layer.
          qkv_bias (bool): If set to False, the qkv layer will not learn an additive
              bias. Default: False.
          dropout_rate (float): Dropout rate.
          kernel_q (_size_3_t): Pooling kernel size for q. If both pooling kernel
              size and pooling stride size are 1 for all the dimensions, pooling is
              disabled.
          kernel_kv (_size_3_t): Pooling kernel size for kv. If both pooling kernel
              size and pooling stride size are 1 for all the dimensions, pooling is
              disabled.
          stride_q (_size_3_t): Pooling kernel stride for q.
          stride_kv (_size_3_t): Pooling kernel stride for kv.
          norm_layer (nn.Module): Normalization layer used after pooling.
          has_cls_embed (bool): If set to True, the first token of the input tensor
              should be a cls token. Otherwise, the input tensor does not contain a
              cls token. Pooling is not applied to the cls token.
          pool_mode (str): Pooling mode. Option includes "conv" (learned pooling), "avg"
              (average pooling), and "max" (max pooling).
          pool_first (bool): If set to True, pool is applied before qkv projection.
              Otherwise, pool is applied after qkv projection. Default: False.
          residual_pool (bool): If set to True, use Improved Multiscale Vision
              Transformer's pooling residual connection.
          depthwise_conv (bool): Whether use depthwise or full convolution for pooling.
          bias_on (bool): Whether use biases for linear layers.
          separate_qkv (bool): Whether to use separate or one layer for qkv projections.
      """
  
      super().__init__()
      assert pool_mode in ["conv", "avg", "max"]
      assert not pool_first
      assert not has_cls_embed
      assert not separate_qkv
      self.dropout_rate = dropout_rate
      self.num_heads = num_heads
      head_dim = dim // num_heads
      self.scale = head_dim**-0.5
      self.has_cls_embed = has_cls_embed
      self.residual_pool = residual_pool
      self.separate_qkv = separate_qkv
      padding_q = [int(q // 2) for q in kernel_q]
      padding_kv = [int(kv // 2) for kv in kernel_kv]
  
      self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
      self.proj = nn.Linear(dim, dim, bias=True if bias_on else False)
  
      # Skip pooling with kernel and stride size of (1, 1, 1).
      if (
          kernel_q is not None
          and numpy.prod(kernel_q) == 1
          and numpy.prod(stride_q) == 1
      ):
          kernel_q = None
      if (
          kernel_kv is not None
          and numpy.prod(kernel_kv) == 1
          and numpy.prod(stride_kv) == 1
      ):
          kernel_kv = None
  
      if pool_mode in ("avg", "max"):
          pool_op = nn.MaxPool3d if pool_mode == "max" else nn.AvgPool3d
          self.pool_q = (
              pool_op(kernel_q, stride_q, padding_q, ceil_mode=False)
              if kernel_q is not None
              else None
          )
          self.pool_k = (
              pool_op(kernel_kv, stride_kv, padding_kv, ceil_mode=False)
              if kernel_kv is not None
              else None
          )
          self.pool_v = (
              pool_op(kernel_kv, stride_kv, padding_kv, ceil_mode=False)
              if kernel_kv is not None
              else None
          )
      elif pool_mode == "conv":
          self.pool_q = (
              nn.Conv3d(
                  head_dim,
                  head_dim,
                  kernel_q,
                  stride=stride_q,
                  padding=padding_q,
                  groups=head_dim if depthwise_conv else 1,
                  bias=False,
              )
              if kernel_q is not None
              else None
          )
          self.pool_k = (
              nn.Conv3d(
                  head_dim,
                  head_dim,
                  kernel_kv,
                  stride=stride_kv,
                  padding=padding_kv,
                  groups=head_dim if depthwise_conv else 1,
                  bias=False,
              )
              if kernel_kv is not None
              else None
          )
          self.pool_v = (
              nn.Conv3d(
                  head_dim,
                  head_dim,
                  kernel_kv,
                  stride=stride_kv,
                  padding=padding_kv,
                  groups=head_dim if depthwise_conv else 1,
                  bias=False,
              )
              if kernel_kv is not None
              else None
          )
      else:
          raise NotImplementedError(f"Unsupported model {pool_mode}")
  
imports: from typing import List, Optional, Tuple
  import numpy
  import torch
  import torch.fx
  import torch.nn
  from torch.nn.common_types import _size_3_t
  from .drop_path import DropPath
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, dropout_rate, kernel_q, kernel_kv, stride_q, stride_kv, norm_layer, has_cls_embed, pool_mode, pool_first, residual_pool, depthwise_conv, bias_on, separate_qkv)