    def __init__(
        self,
        pool: Optional[Module],
        has_cls_embed: bool,
        norm: Optional[Module],
    ) -> None:
        """Apply pool to a flattened input (given pool operation and the unflattened shape).


                                         Input
                                           ↓
                                        Reshape
                                           ↓
                                          Pool
                                           ↓
                                        Reshape
                                           ↓
                                          Norm


        Params:
            pool (Optional[Callable]): Pool operation that is applied to the input tensor.
                If pool is none, return the input tensor.
            has_cls_embed (bool): Whether the input tensor contains cls token. Pool
                operation excludes cls token.
            norm: (Optional[Callable]): Optional normalization operation applied to
            tensor after pool.
        """
        super().__init__()
        self.has_pool = pool is not None
        self.pool = pool if pool is not None else Identity()

        self.has_cls_embed = has_cls_embed
        if norm is not None:
            self.norm_before_pool = isinstance(norm, (BatchNorm3d, Identity))
            self.has_norm = True
            self.norm = norm
        else:
            self.norm_before_pool = False
            self.has_norm = False
            self.norm = Identity
Apply pool to a flattened input (given pool operation and the unflattened shape).


Input
↓
Reshape
↓
Pool
↓
Reshape
↓
Norm


Params:
pool (Optional[Callable]): Pool operation that is applied to the input tensor.
If pool is none, return the input tensor.
has_cls_embed (bool): Whether the input tensor contains cls token. Pool
operation excludes cls token.
norm: (Optional[Callable]): Optional normalization operation applied to
tensor after pool.
python/aitemplate/frontend/nn/multiscale_attention.py