    def __init__(
        self,
        *,
        patch_embed: Optional[Module],
        cls_positional_encoding: Module,
        pos_drop: Optional[Module],
        blocks: ModuleList,
        norm_embed: Optional[Module],
        head: Optional[Module],
    ) -> None:
        """
        Args:
            patch_embed (nn.Module): Patch embed module.
            cls_positional_encoding (nn.Module): Positional encoding module.
            pos_drop (Optional[nn.Module]): Dropout module after patch embed.
            blocks (nn.ModuleList): Stack of multi-scale transformer blocks.
            norm_layer (nn.Module): Normalization layer before head.
            head (Optional[nn.Module]): Head module.
        """
        super().__init__()

        assert hasattr(
            cls_positional_encoding, "patch_embed_shape"
        ), "cls_positional_encoding should have method patch_embed_shape."

        self.patch_embed = patch_embed or Identity()
        self.cls_positional_encoding = cls_positional_encoding
        self.pos_drop = pos_drop or Identity()
        self.blocks = blocks
        self.norm_embed = norm_embed or Identity()
        self.head = head or Identity()
Args:
patch_embed (nn.Module): Patch embed module.
cls_positional_encoding (nn.Module): Positional encoding module.
pos_drop (Optional[nn.Module]): Dropout module after patch embed.
blocks (nn.ModuleList): Stack of multi-scale transformer blocks.
norm_layer (nn.Module): Normalization layer before head.
head (Optional[nn.Module]): Head module.
python/aitemplate/frontend/nn/vision_transformers.py