    def __init__(
        self,
        dim: int,
        dim_out: int,
        num_heads: int,
        seq_len: int,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = False,
        dropout_rate: float = 0.0,
        droppath_rate: float = 0.0,
        act_layer: Module = GELU,
        norm_layer: Module = LayerNorm,
        attn_norm_layer: Module = LayerNorm,
        kernel_q=(1, 1, 1),
        kernel_kv=(1, 1, 1),
        stride_q=(1, 1, 1),
        stride_kv=(1, 1, 1),
        pool_mode: str = "conv",
        has_cls_embed: bool = True,
        pool_first: bool = False,
        residual_pool: bool = False,
        depthwise_conv: bool = True,
        bias_on: bool = True,
        separate_qkv: bool = False,
    ) -> None:
        """
        Args:
            dim (int): Input feature dimension.
            dim_out (int): Output feature dimension.
            num_heads (int): Number of heads in the attention layer.
            mlp_ratio (float): Mlp ratio which controls the feature dimension in the
                hidden layer of the Mlp block.
            qkv_bias (bool): If set to False, the qkv layer will not learn an additive
                bias. Default: False.
            dropout_rate (float): DropOut rate. If set to 0, DropOut is disabled.
            droppath_rate (float): DropPath rate. If set to 0, DropPath is disabled.
            act_layer (Module): Activation layer used in the Mlp layer.
            norm_layer (Module): Normalization layer.
            attn_norm_layer (Module): Normalization layer in the attention module.
            kernel_q (_size_3_t): Pooling kernel size for q. If pooling kernel size is
                1 for all the dimensions, pooling is not used (by default).
            kernel_kv (_size_3_t): Pooling kernel size for kv. If pooling kernel size
                is 1 for all the dimensions, pooling is not used. By default, pooling
                is disabled.
            stride_q (_size_3_t): Pooling kernel stride for q.
            stride_kv (_size_3_t): Pooling kernel stride for kv.
            pool_mode (str): Pooling mode. Option includes "conv" (learned pooling), "avg"
                (average pooling), and "max" (max 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_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__()
        self.dim = dim
        self.dim_out = dim_out
        self.norm1 = norm_layer(dim)
        self.norm1_is_batchnorm_1d = isinstance(self.norm1, BatchNorm1d)
        self.norm1.permute_input_output = True if self.norm1_is_batchnorm_1d else False
        kernel_skip = [s + 1 if s > 1 else s for s in stride_q]
        stride_skip = stride_q
        padding_skip = [int(skip // 2) for skip in kernel_skip]
        self.attn = MultiScaleAttention(
            dim,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            dropout_rate=dropout_rate,
            kernel_q=kernel_q,
            kernel_kv=kernel_kv,
            stride_q=stride_q,
            stride_kv=stride_kv,
            norm_layer=attn_norm_layer,
            has_cls_embed=has_cls_embed,
            pool_mode=pool_mode,
            pool_first=pool_first,
            residual_pool=residual_pool,
            bias_on=bias_on,
            depthwise_conv=depthwise_conv,
            separate_qkv=separate_qkv,
            max_seq_len=seq_len,
        )
        self.drop_path = DropPath(droppath_rate) if droppath_rate > 0.0 else Identity()
        self.norm2 = norm_layer(dim)
        self.norm2_is_batchnorm_1d = isinstance(self.norm2, BatchNorm1d)
        self.norm2.permute_input_output = True if self.norm2_is_batchnorm_1d else False
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.has_cls_embed = has_cls_embed
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            out_features=dim_out,
            act_layer=act_layer,
            dropout_rate=dropout_rate,
            bias_on=bias_on,
        )
        if dim != dim_out:
            self.proj = Linear(dim, dim_out, bias=bias_on)
        else:
            self.proj = Identity()

        self.pool_skip = (
            MaxPool3d(tuple(kernel_skip), tuple(stride_skip), tuple(padding_skip))
            if len(stride_skip) > 0 and numpy.prod(stride_skip) > 1
            else None
        )
        self._attention_pool = _AttentionPool(
            self.pool_skip, has_cls_embed=self.has_cls_embed, norm=None
        )
Args:
dim (int): Input feature dimension.
dim_out (int): Output feature dimension.
num_heads (int): Number of heads in the attention layer.
mlp_ratio (float): Mlp ratio which controls the feature dimension in the
hidden layer of the Mlp block.
qkv_bias (bool): If set to False, the qkv layer will not learn an additive
bias. Default: False.
dropout_rate (float): DropOut rate. If set to 0, DropOut is disabled.
droppath_rate (float): DropPath rate. If set to 0, DropPath is disabled.
act_layer (Module): Activation layer used in the Mlp layer.
norm_layer (Module): Normalization layer.
attn_norm_layer (Module): Normalization layer in the attention module.
kernel_q (_size_3_t): Pooling kernel size for q. If pooling kernel size is
1 for all the dimensions, pooling is not used (by default).
kernel_kv (_size_3_t): Pooling kernel size for kv. If pooling kernel size
is 1 for all the dimensions, pooling is not used. By default, pooling
is disabled.
stride_q (_size_3_t): Pooling kernel stride for q.
stride_kv (_size_3_t): Pooling kernel stride for kv.
pool_mode (str): Pooling mode. Option includes "conv" (learned pooling), "avg"
(average pooling), and "max" (max 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_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.
python/aitemplate/frontend/nn/multiscale_attention.py