    def __init__(
        self,
        dim: int,
        num_heads: int = 8,
        batch_size: int = 1,
        qkv_bias: bool = False,
        dropout_rate: float = 0.0,
        kernel_q=(1, 1, 1),
        kernel_kv=(1, 1, 1),
        stride_q=(1, 1, 1),
        stride_kv=(1, 1, 1),
        norm_layer: Callable = 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 = False,
        max_seq_len: int = 6272,
    ) -> 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 (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"]

        self.pool_first = pool_first
        self.dropout_rate = dropout_rate
        self.num_heads = num_heads
        self.dim = dim
        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
        self.max_seq_len = max_seq_len
        padding_q = [int(q // 2) for q in kernel_q]
        padding_kv = [int(kv // 2) for kv in kernel_kv]

        # Set placeholders for torchscriptability, may not be actually used
        self.q = self.k = self.v = self.qkv = Identity()
        if self.pool_first or self.separate_qkv:
            self.q = Linear(dim, dim, bias=qkv_bias)
            self.k = Linear(dim, dim, bias=qkv_bias)
            self.v = Linear(dim, dim, bias=qkv_bias)
        else:
            self.qkv = Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = Linear(dim, dim, bias=True if bias_on else False)

        assert dropout_rate == 0.0
        if dropout_rate > 0.0:
            self.proj_drop = Dropout(dropout_rate)
        else:
            self.proj_drop = Identity()

        # Skip pooling with kernel and stride size of (1, 1, 1).
        if (
            kernel_q is not None
            and self._prod(kernel_q) == 1
            and self._prod(stride_q) == 1
        ):
            kernel_q = None
        if (
            kernel_kv is not None
            and self._prod(kernel_kv) == 1
            and self._prod(stride_kv) == 1
        ):
            kernel_kv = None

        if pool_mode in ["max", "avg"]:
            raise NotImplementedError(f"Unsupported input dimension {pool_mode}")

        ## TODO: add pool mode support for {"max", "avg"}

        elif pool_mode == "conv":
            self.pool_q = (
                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.norm_q = norm_layer(head_dim) if kernel_q is not None else None
            self.pool_k = (
                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.norm_k = norm_layer(head_dim) if kernel_kv is not None else None
            self.pool_v = (
                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.norm_v = norm_layer(head_dim) if kernel_kv is not None else None
        else:
            raise NotImplementedError(f"Unsupported model {pool_mode}")

        # Will not be used if `separate_qkv == True`
        self._attention_pool_q = _AttentionPool(
            self.pool_q,
            has_cls_embed=self.has_cls_embed,
            norm=getattr(self, "norm_q", None),
        )
        self._attention_pool_k = _AttentionPool(
            self.pool_k,
            has_cls_embed=self.has_cls_embed,
            norm=getattr(self, "norm_k", None),
        )
        self._attention_pool_v = _AttentionPool(
            self.pool_v,
            has_cls_embed=self.has_cls_embed,
            norm=getattr(self, "norm_v", 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 (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.
python/aitemplate/frontend/nn/multiscale_attention.py