def __init__(
        self,
        img_size=224,
        patch_size=16,
        in_chans=3,
        num_classes=1000,
        d_model=768,
        depth=12,
        # num_heads=12,
        expand=4,
        # qkv_bias=True,
        # qk_scale=None,
        representation_size=None,
        distilled=False,
        dropout=0.,
        # attn_drop_rate=0.,
        drop_path_rate=0.,
        embed_layer=PatchEmbed,
        norm='layer',
        # norm_layer=None,
        # act_layer=None,
        weight_init='',
        # attnlinear_cfg=None,
        # mlp_cfg=None,
        layer=None,
        # ff_cfg=None,
        transposed=False,
        layer_reps=1,
        use_pos_embed=False,
        use_cls_token=False,
        track_norms=False,
    ):
        """
        Args:
            img_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_chans (int): number of input channels
            num_classes (int): number of classes for classification head
            d_model (int): embedding dimension
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            distilled (bool): model includes a distillation token and head as in DeiT models
            dropout (float): dropout rate
            attn_drop_rate (float): attention dropout rate
            drop_path_rate (float): stochastic depth rate
            embed_layer (nn.Module): patch embedding layer
            norm_layer: (nn.Module): normalization layer
            weight_init: (str): weight init scheme
        """
        super().__init__()
        self.num_classes = num_classes
        self.num_features = self.d_model = d_model  # num_features for consistency with other models
        self.num_tokens = 2 if distilled else 1
        self.use_pos_embed = use_pos_embed
        self.use_cls_token = use_cls_token
        # norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
        # act_layer = act_layer or nn.GELU

        self.track_norms = track_norms

        self.patch_embed = embed_layer(
            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=d_model,
        )
        num_patches = self.patch_embed.num_patches

        self.cls_token = None
        self.dist_token = None
        if use_cls_token:
            self.cls_token = nn.Parameter(torch.zeros(1, 1, d_model))
            self.dist_token = nn.Parameter(torch.zeros(1, 1, d_model)) if distilled else None
        else:
            assert not distilled, 'Distillation token not supported without class token'

        self.pos_embed = None
        if use_pos_embed:
            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, d_model))
            self.pos_drop = nn.Dropout(p=dropout)

        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
        # self.blocks = nn.Sequential(*[
        #     Block(
        #         dim=d_model, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
        #         drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, act_layer=act_layer,
        #         attnlinear_cfg=attnlinear_cfg, mlp_cfg=mlp_cfg)
        #     for i in range(depth)
        # ])

        self.transposed = transposed

        layer = to_list(layer, recursive=False) * layer_reps

        # Some special arguments are passed into each layer
        for _layer in layer:
            # If layers don't specify dropout, add it
            if _layer.get('dropout', None) is None:
                _layer['dropout'] = dropout
            # Ensure all layers are shaped the same way
            _layer['transposed'] = transposed

        # # Layer arguments
        # layer_cfg = layer.copy()
        # layer_cfg['dropout'] = dropout
        # layer_cfg['transposed'] = self.transposed
        # layer_cfg['initializer'] = None
        # # layer_cfg['l_max'] = L
        # print("layer config", layer_cfg)

        # Config for the inverted bottleneck
        ff_cfg = {
            '_name_': 'ffn',
            'expand': int(expand),
            'transposed': self.transposed,
            'activation': 'gelu',
            'initializer': None,
            'dropout': dropout,
        }

        blocks = []
        for i in range(depth):
            for _layer in layer:
                blocks.append(
                    SequenceResidualBlock(
                        d_input=d_model,
                        i_layer=i,
                        prenorm=True,
                        dropout=dropout,
                        layer=_layer,
                        residual='R',
                        norm=norm,
                        pool=None,
                        drop_path=dpr[i],
                    )
                )
            if expand > 0:
                blocks.append(
                    SequenceResidualBlock(
                        d_input=d_model,
                        i_layer=i,
                        prenorm=True,
                        dropout=dropout,
                        layer=ff_cfg,
                        residual='R',
                        norm=norm,
                        pool=None,
                        drop_path=dpr[i],
                    )
                )
        self.blocks = nn.Sequential(*blocks)

        # self.norm = norm_layer(d_model)
        if norm is None:
            self.norm = None
        elif isinstance(norm, str):
            self.norm = Normalization(d_model, transposed=self.transposed, _name_=norm)
        else:
            self.norm = Normalization(d_model, transposed=self.transposed, **norm)

        # Representation layer: generally defaults to nn.Identity()
        if representation_size and not distilled:
            self.num_features = representation_size
            self.pre_logits = nn.Sequential(OrderedDict([
                ('fc', nn.Linear(d_model, representation_size)),
                ('act', nn.Tanh())
            ]))
        else:
            self.pre_logits = nn.Identity()

        # Classifier head(s): TODO: move to decoder
        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
        self.head_dist = None
        if distilled:
            self.head_dist = nn.Linear(self.d_model, self.num_classes) if num_classes > 0 else nn.Identity()

        # Weight init
        assert weight_init in ('jax', 'jax_nlhb', 'nlhb', '')
        head_bias = -math.log(self.num_classes) if 'nlhb' in weight_init else 0.
        if self.pos_embed is not None:
            trunc_normal_(self.pos_embed, std=.02)
        if self.dist_token is not None:
            trunc_normal_(self.dist_token, std=.02)
        if weight_init.startswith('jax'):
            # leave cls token as zeros to match jax impl
            for n, m in self.named_modules():
                _init_vit_weights(m, n, head_bias=head_bias, jax_impl=True)
        else:
            if self.cls_token is not None:
                trunc_normal_(self.cls_token, std=.02)
            self.apply(_init_vit_weights)

        Args:
            img_size (int, tuple): input image size
            patch_size (int, tuple): patch size
            in_chans (int): number of input channels
            num_classes (int): number of classes for classification head
            d_model (int): embedding dimension
            depth (int): depth of transformer
            num_heads (int): number of attention heads
            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
            qkv_bias (bool): enable bias for qkv if True
            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
            representation_size (Optional[int]): enable and set representation layer (pre-logits) to this value if set
            distilled (bool): model includes a distillation token and head as in DeiT models
            dropout (float): dropout rate
            attn_drop_rate (float): attention dropout rate
            drop_path_rate (float): stochastic depth rate
            embed_layer (nn.Module): patch embedding layer
            norm_layer: (nn.Module): normalization layer
            weight_init: (str): weight init scheme
        
analyze/convnexts4nd/vit_all.py