def vit_base_s4nd():
    """
    model:
    _name_: vit_b_16
    dropout: 0.0
    drop_path_rate: 0.1
    d_model: 768
    depth: 12
    expand: 4
    norm: layer
    layer_reps: 1
    use_cls_token: false
    use_pos_embed: false

    layer:
        d_state: 64
        final_act: glu
        bidirectional: true
        channels: 2
        lr: 0.001
        n_ssm: 1
        contract_version: 1  # 0 is for 2d, 1 for 1d or 3d (or other)
    """
    """
    defaults:
    - layer: vit

    patch_size: 16
    d_model: 768
    dropout: 0.0
    drop_path_rate: 0.0
    depth: 8
    expand: 3
    norm: layer
    use_pos_embed: true
    use_cls_token: true
    """

    model = VisionTransformer(
        dropout=0.0,
        drop_path_rate=0.1,
        d_model=768,
        depth=12,
        expand=4,
        norm="layer",
        layer_reps=1,
        use_cls_token=False,
        use_pos_embed=False,
        layer=dict(
            _name_="s4nd",
            d_state=64,
            final_act="glu",
            bidirectional=True,
            channels=2,
            lr=0.001,
            n_ssm=1,
            contract_version=1,
            # contract_version=0,
            activation="gelu",
            initializer=None,
            weight_norm=False,
            trank=1,
            dropout=0,
            tie_dropout=0,
            init="legs",
            dt_min=0.001,
            dt_max=0.1,
            deterministic=False,
            l_max=None,
            verbose=True,
            linear=False,
        )
    )
    
    def forward(self, x, rate=1.0, resolution=None, state=None):
        x = self.forward_features(x)
        if self.head_dist is not None:
            x, x_dist = self.head(x[0]), self.head_dist(x[1])  # x must be a tuple
            if self.training and not torch.jit.is_scripting():
                # during inference, return the average of both classifier predictions
                return x, x_dist
            else:
                return (x + x_dist) / 2
        else:
            x = self.head(x)
        return x
    
    model.forward = partial(forward, model)
    
    return model

    model:
    _name_: vit_b_16
    dropout: 0.0
    drop_path_rate: 0.1
    d_model: 768
    depth: 12
    expand: 4
    norm: layer
    layer_reps: 1
    use_cls_token: false
    use_pos_embed: false

    layer:
        d_state: 64
        final_act: glu
        bidirectional: true
        channels: 2
        lr: 0.001
        n_ssm: 1
        contract_version: 1  # 0 is for 2d, 1 for 1d or 3d (or other)
    
analyze/convnexts4nd/vit_all.py