Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1 @@
+fabric*

config.json

@@ -0,0 +1,50 @@
+{
+  "_name_or_path": "mistralai/mamba-codestral-7B-v0.1",
+  "architectures": [
+    "Mamba2ForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_ibs2.IBS2Config",
+    "AutoModel": "modeling_ibs2.IBS2ForCausalLM",
+    "AutoModelForCausalLM": "modeling_ibs2.IBS2ForCausalLM"
+  },
+  "bos_token_id": 0,
+  "chunk_size": 256,
+  "conv_kernel": 4,
+  "eos_token_id": 0,
+  "expand": 2,
+  "head_dim": 64,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "ib_type": "gamma",
+  "initializer_range": 0.1,
+  "intermediate_size": 8192,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "ibs2",
+  "n_groups": 8,
+  "norm_before_gate": true,
+  "num_heads": 128,
+  "num_hidden_layers": 64,
+  "pad_token_id": 0,
+  "rescale_prenorm_residual": false,
+  "residual_in_fp32": true,
+  "rms_norm": true,
+  "state_size": 128,
+  "tie_word_embeddings": false,
+  "time_step_floor": 0.0001,
+  "time_step_init_scheme": "random",
+  "time_step_limit": [
+    0.0,
+    Infinity
+  ],
+  "time_step_max": 0.1,
+  "time_step_min": 0.001,
+  "time_step_rank": 256,
+  "time_step_scale": 1.0,
+  "torch_dtype": "float32",
+  "transformers_version": "4.43.3",
+  "use_bias": false,
+  "use_cache": true,
+  "use_conv_bias": true,
+  "vocab_size": 32768
+}

configuration_ibs2.py

@@ -0,0 +1,189 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MAMBA2 configuration"""
+
+import math
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class IBS2Config(PretrainedConfig):
+    """
+    This is the configuration class to store the configuration of a [`Mamba2Model`]. It is used to instantiate a MAMBA2
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to that of the MAMBA2
+    [state-spaces/mamba2-2.8b](https://huggingface.co/state-spaces/mamba2-2.8b) architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        num_heads (`int`, *optional*, defaults to 128):
+            Number of heads for the evolution matrices of mamba 2.
+        head_dim (`int`, *optional*, defaults to 64):
+            Dimension of each head.
+        vocab_size (`int`, *optional*, defaults to 32768):
+            Vocabulary size of the MAMBA2 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`Mamba2Model`].
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimensionality of the embeddings and hidden states.
+        state_size (`int`, *optional*, defaults to 128): shape of the state space latents.
+        num_hidden_layers (`int`, *optional*, defaults to 64):
+            Number of hidden layers in the model.
+        layer_norm_epsilon (`float`, *optional*, defaults to 1e-05):
+            The epsilon to use in the layer normalization layers.
+        pad_token_id (`int`, *optional*, defaults to 1):
+            Padding token id.
+        bos_token_id (`int`, *optional*, defaults to 0):
+            The id of the beginning of sentence token in the vocabulary.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            The id of the end of sentence token in the vocabulary.
+        expand (`int`, *optional*, defaults to 2): Expanding factor used to determine the intermediate size.
+        conv_kernel (`int`, *optional*, defaults to 4): Size of the convolution kernel.
+        n_groups (`int`, *optional*, defaults to 8):
+            Number of groups for the evolution matrices of mamba 2.
+        use_bias (`bool`, *optional*, defaults to `False`):
+            Whether or not to use bias in ["in_proj", "out_proj"] of the mixer block
+        use_conv_bias (`bool`, *optional*, defaults to `True`):
+            Whether or not to use bias in the convolution layer of the mixer block.
+        hidden_act (`str`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        initializer_range (`float`, *optional*, defaults to 0.1):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        residual_in_fp32 (`bool`, *optional*, defaults to `True`):
+            Whether or not residuals should be in `float32`. If set to `False` residuals will keep the same `dtype` as the rest of the model
+        time_step_rank (`Union[int,str]`, *optional*, defaults to `"auto"`):
+            Rank of the discretization projection matrix. `"auto"` means that it will default to `math.ceil(self.hidden_size / 16)`
+        time_step_min (`float`, *optional*, defaults to 0.001):
+            Minimum `time_step` used to bound `dt_proj.bias`.
+        time_step_max (`float`, *optional*, defaults to 0.1):
+            Maximum `time_step` used to bound `dt_proj.bias`.
+        time_step_floor (`float`, *optional*, defaults to 0.0001):
+            Minimum clamping value of the `dt_proj.bias` layer initialization.
+        time_step_limit (`tuple`, *optional*, defaults to `(0.0, inf)`):
+            Accepted range of time step values.
+        rescale_prenorm_residual (`bool`, *optional*, defaults to `False`):
+            Whether or not to rescale `out_proj` weights when initializing.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the cache should be used.
+        rms_norm (`bool`, *optional*, defaults to `True`):
+            Whether to use RMS norm or not.
+        chunk_size (`int`, *optional*, defaults to 256):
+            Size of the chunks that will comprise the sequence.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether to tie word embeddings or not.
+
+
+    Example:
+
+    ```python
+    >>> from transformers import Mamba2Config, Mamba2Model
+
+    >>> # Initializing a Mamba2 configuration
+    >>> configuration = Mamba2Config()
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = Mamba2Model(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "ibs2"
+
+    def __init__(
+        self,
+        num_classes=1,
+        ib_type=None,
+        return_attn=False,
+        num_heads=128,
+        head_dim=64,
+        vocab_size=32768,
+        hidden_size=4096,
+        state_size=128,
+        num_hidden_layers=64,
+        layer_norm_epsilon=1e-5,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=2,
+        expand=2,
+        conv_kernel=4,
+        n_groups=8,
+        use_bias=False,
+        use_conv_bias=True,
+        hidden_act="silu",
+        initializer_range=0.1,
+        residual_in_fp32=True,
+        time_step_rank="auto",
+        time_step_min=0.001,
+        time_step_max=0.1,
+        time_step_floor=1e-4,
+        time_step_limit=(0.0, float("inf")),
+        rescale_prenorm_residual=False,
+        use_cache=True,
+        rms_norm=True,
+        chunk_size=256,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.num_classes = num_classes
+        self.ib_type = ib_type
+        self.return_attn = return_attn
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.state_size = state_size
+        self.num_hidden_layers = num_hidden_layers
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.conv_kernel = conv_kernel
+        self.expand = expand
+
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.pad_token_id = pad_token_id
+        self.use_bias = use_bias
+        self.use_conv_bias = use_conv_bias
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.time_step_rank = math.ceil(self.hidden_size / 16) if time_step_rank == "auto" else time_step_rank
+        self.time_step_min = time_step_min
+        self.time_step_max = time_step_max
+        self.time_step_floor = time_step_floor
+        self.rescale_prenorm_residual = rescale_prenorm_residual
+        self.residual_in_fp32 = residual_in_fp32
+        self.use_cache = use_cache
+        self.n_groups = n_groups
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.rms_norm = rms_norm
+        self.state_size = state_size
+        self.chunk_size = chunk_size
+        self.time_step_limit = time_step_limit
+        self.tie_word_embeddings = tie_word_embeddings
+
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["IBS2Config"]

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "pad_token_id": 1,
+  "transformers_version": "4.43.3"
+}

modeling_ibs2.py

@@ -0,0 +1,1457 @@
+# coding=utf-8
+# Copyright 2024 state-spaces/mamba2 org and HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch MAMBA2 model."""
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from transformers.activations import ACT2FN
+from transformers.generation import GenerationMixin
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    ModelOutput,
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+)
+from transformers.utils.import_utils import is_causal_conv1d_available, is_torch_available, _is_package_available, version
+from .configuration_ibs2 import IBS2Config
+
+def is_mamba_2_ssm_available():
+    if is_torch_available():
+        import torch
+
+        if not torch.cuda.is_available():
+            return False
+        else:
+            if _is_package_available("mamba_ssm"):
+                import mamba_ssm
+
+                if version.parse(mamba_ssm.__version__) >= version.parse("2.0.4"):
+                    return True
+    return False
+
+logger = logging.get_logger(__name__)
+
+
+if is_mamba_2_ssm_available():
+    from mamba_ssm.ops.triton.selective_state_update import selective_state_update
+    from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
+else:
+    mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined, selective_state_update = None, None, None
+
+if is_causal_conv1d_available():
+    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+else:
+    causal_conv1d_update, causal_conv1d_fn = None, None
+
+is_fast_path_available = all(
+    (
+        selective_state_update,
+        mamba_chunk_scan_combined,
+        mamba_split_conv1d_scan_combined,
+        causal_conv1d_fn,
+        causal_conv1d_update,
+    )
+)
+
+_CHECKPOINT_FOR_DOC = "mistralai/mamba-codestral-7B-v0.1"
+_CONFIG_FOR_DOC = "Mamba2Config"
+
+
+# Helper methods for segment sum computation
+
+
+def pad_tensor_by_size(input_tensor: torch.Tensor, pad_size: int):
+    """
+    Padding x tensor with `pad_size` on the seq_len dim (dim=1)
+
+    Assumes that we only have tensors of either size 4 or 3
+    """
+    pad_shape = (0, 0, 0, 0, 0, pad_size, 0, 0) if len(input_tensor.shape) == 4 else (0, 0, 0, pad_size, 0, 0)
+
+    return torch.nn.functional.pad(input_tensor, pad_shape, mode="constant", value=0)
+
+
+def reshape_into_chunks(input_tensor, pad_size, chunk_size):
+    """
+    Padding input_tensor with `pad_size` on the seq_len dim (dim=1) and
+    simultaneously splitting it into chunk sequences.
+
+    Assumes that we only have tensors of either size 4 or 3
+    """
+    # [bsz, seq_len, ...] -> [bsz, seq_len multiple of chunk_size, ...]
+    input_tensor = pad_tensor_by_size(input_tensor, pad_size)
+
+    if len(input_tensor.shape) == 3:
+        # [bsz, seq_len multiple of chunk_size, num_heads] -> [bsz, -1, chunk_size, num_heads]
+        return input_tensor.reshape(input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2])
+    else:
+        # [bsz, seq_len multiple of chunk_size, num_heads, head_dim or state_size] -> [bsz, -1, chunk_size, num_heads, head_dim or state_size]
+        return input_tensor.reshape(
+            input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2], input_tensor.shape[3]
+        )
+
+
+def segment_sum(input_tensor):
+    """
+    More stable segment sum calculation. Uses cumulative sums and masking instead of direct subtractions.
+    """
+    chunk_size = input_tensor.size(-1)
+    # 1. expand input tensor to have an additional dimension and repeat along that dimension
+    # [..., chunk_size] -> [..., chunk_size, chunk_size]
+    input_tensor = input_tensor[..., None].expand(*input_tensor.size(), chunk_size)
+    # 2. create a lower triangular mask with the diagonal set to 0 to 0 out elements above diag
+    mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=-1)
+    input_tensor = input_tensor.masked_fill(~mask, 0)
+    # 3. compute actual cumsum
+    tensor_segsum = torch.cumsum(input_tensor, dim=-2)
+
+    # 4. apply mask to keep only the lower triangular part of the cumulative sum result (incl diagonal this time)
+    mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=0)
+    tensor_segsum = tensor_segsum.masked_fill(~mask, -torch.inf)
+    return tensor_segsum
+
+
+def apply_mask_to_padding_states(hidden_states, attention_mask):
+    """
+    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
+    """
+    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+        dtype = hidden_states.dtype
+        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+    return hidden_states
+
+
+class Mamba2Cache:
+    """
+    Arguments:
+        config: Mamba2Config
+        batch_size: int
+        dtype: torch.dtype
+        device: torch.device
+
+    Attributes:
+        dtype: (`torch.dtype`):
+            The default `dtype` used to initializing the cache.
+        conv_kernel_size: (`int`):
+            Model's convolution kernel size taken from config.
+        n_groups: (`int`):
+            Model's number of groups taken from the config - similar to tensor parallel in Transformer.
+        state_size: (`int`):
+            Model's SSM state size taken from config.
+        num_heads: (`int`):
+            The number of heads used in the linear attention / SSM.
+        head_dim: (`int`):
+            The respective dimension of the heads used in the linear attention / SSM.
+        intermediate_size: (`int`):
+            Model's intermediate_size based on (expand * hidden_dim) from config.
+        conv_states: (`torch.Tensor`):
+            A tensor of shape `[num_layers, batch_size, conv_kernel_size, intermediate_size + 2 * n_groups * state_size]` that holds convolutional states.
+        ssm_states: (`torch.Tensor`):
+            A tensor of shape `[num_layers, batch_size, num_heads, head_dim, state_size]` that holds ssm states.
+    """
+
+    def __init__(
+        self, config: IBS2Config, batch_size: int, dtype: torch.dtype = torch.float16, device: Optional[str] = None
+    ):
+        self.dtype = dtype
+        self.conv_kernel_size = config.conv_kernel
+        self.n_groups = config.n_groups
+        self.state_size = config.state_size
+        self.num_heads = config.num_heads
+        self.head_dim = config.head_dim
+        self.intermediate_size = int(config.expand * config.hidden_size)
+
+        self.conv_states = torch.zeros(
+            config.num_hidden_layers,
+            batch_size,
+            self.intermediate_size + 2 * self.n_groups * self.state_size,
+            self.conv_kernel_size,
+            device=device,
+            dtype=dtype,
+        )
+        self.ssm_states = torch.zeros(
+            config.num_hidden_layers,
+            batch_size,
+            self.num_heads,
+            self.head_dim,
+            self.state_size,
+            device=device,
+            dtype=dtype,
+        )
+
+    def update_conv_state(
+        self, layer_idx: int, new_conv_state: torch.Tensor, cache_init: bool = False
+    ) -> torch.Tensor:
+        if cache_init:
+            self.conv_states[layer_idx] = new_conv_state.to(self.conv_states.device)
+        else:
+            self.conv_states[layer_idx] = self.conv_states[layer_idx].roll(shifts=-1, dims=-1)
+            self.conv_states[layer_idx][:, :, -1] = new_conv_state[:, 0, :].to(self.conv_states.device)
+        return self.conv_states[layer_idx]
+
+    def update_ssm_state(self, layer_idx: int, new_ssm_state: torch.Tensor):
+        self.ssm_states[layer_idx] = new_ssm_state.to(self.ssm_states.device)
+        return self.ssm_states[layer_idx]
+
+    def reset(self):
+        self.conv_states.zero_()
+        self.ssm_states.zero_()
+
+
+class MambaRMSNormGated(torch.nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states, gate=None):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+
+        if gate is not None:
+            hidden_states = hidden_states * nn.functional.silu(gate.to(torch.float32))
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+        return self.weight * hidden_states.to(input_dtype)
+
+class Normalize(nn.Module):
+    def __init__(self, min_value=None, max_value=None):
+        super().__init__()
+        self.min_value = min_value
+        self.max_value = max_value
+
+    def forward(self, value_states):
+        # 计算 value_states 的最小值和最大值
+        min_val = value_states.min(dim=-1, keepdim=True).values
+        max_val = value_states.max(dim=-1, keepdim=True).values
+        
+        # 如果 min_value 或 max_value 是 None，则使用另一个值进行归一化
+        if self.min_value is None and self.max_value is not None:
+            # 只使用 max_value 进行归一化
+            scale_factor = self.max_value / (max_val + 1e-6)
+            return value_states * scale_factor, scale_factor, None
+        elif self.max_value is None and self.min_value is not None:
+            # 只使用 min_value 进行归一化
+            scale_factor = self.min_value / (min_val + 1e-6)
+            return value_states * scale_factor, scale_factor, None
+        elif self.min_value is not None and self.max_value is not None:
+            # 同时使用 min_value 和 max_value 进行归一化
+            scale_factor = (self.max_value - self.min_value) / (max_val - min_val + 1e-6)
+            shift_factor = self.min_value - min_val * scale_factor
+            normalized_value_states = value_states * scale_factor + shift_factor
+            return normalized_value_states, scale_factor, shift_factor
+        else:
+            # 如果 min_value 和 max_value 都是 None，则不进行归一化
+            return value_states, None, None
+
+
+# torch.lgamma(n) - torch.lgamma(theta) + (theta - n) * torch.special.digamma(theta)
+class GammaIB(nn.Module):
+    def __init__(self, hidden_size, alphas=None, return_attn=False, **kwargs) -> None:
+        super().__init__()
+        self.alphas = alphas
+        # self.attributor = nn.Linear(hidden_size, 1)
+        self.hidden_size = hidden_size
+        self._auxiliary_loss = 0
+        self.epoch_frac = 0
+        self.epoch_threshold = -1
+        self.normalizer = Normalize(max_value=10, min_value=0.1)
+        self.return_attn = return_attn
+        self._attn = None
+
+
+
+    def get_auxiliary_loss(self):
+        loss = self._auxiliary_loss
+        self._auxiliary_loss = 0.0
+        # print("auxiliary_loss", loss)
+        return loss
+
+    def init_alphas(self, param_alphas):
+        # shape: [bsz, seq_len, 1]
+        if self.alphas is None:
+            maxmimum = 8
+        else:
+            maxmimum = param_alphas.size(1) / self.alphas.size(0) * self.alphas.max().item()
+        length = param_alphas.shape[1]
+        alphas = torch.linspace(maxmimum, 1, steps=length).float().to(param_alphas.device) # distance-decay: torch.linspace(1, 2, steps=length)
+        self.alphas = alphas
+
+    def compute_loss(self, param_alphas, epsilon=1e-6):
+        if self.alphas is None:
+            self.init_alphas(param_alphas)  # length is the second dimension of att
+            print(f"Gamma prior alpha first: {self.alphas[0]}, last: {self.alphas[-1]}, size: {self.alphas.size(0)}", )
+        if self.alphas.size(0) != param_alphas.size(1):
+            self.init_alphas(param_alphas)  # length is the second dimension of att
+            print(f"Gamma prior alpha first: {self.alphas[0]}, last: {self.alphas[-1]}, size: {self.alphas.size(0)}", )
+        
+        params = self.alphas.unsqueeze(-1).expand(param_alphas.shape)
+        reg_loss = (torch.lgamma(self.alphas) - torch.lgamma(params) + (params - self.alphas) * torch.digamma(params)).mean()
+        
+        return reg_loss
+
+    def forward(self, states):
+        # hidden_states shape [bsz, seq_length, dimension]
+        hidden_states, alphas = torch.split(
+                    states,
+                    [self.hidden_size, 1],
+                    dim=-1
+                )
+        if self.epoch_frac < self.epoch_threshold:
+            return hidden_states
+        # alphas = F.rms_norm(alphas, [hidden_states.shape[1], 1])
+        #! HACK: for finetuning we zero-init and reparametrize to ensure the initialization is one
+        alphas = nn.functional.softplus(alphas) - torch.log(torch.tensor(2.0)) + torch.tensor(1.0)
+        # if self.return_attn:
+        #     # print(alphas.shape)
+        #     if alphas.shape[-2] == 1:
+        #         pass # seqlen == 1 indicates it is caching
+        #     else:
+        #         self._attn = alphas.detach().cpu()
+
+        shaped_alphas = alphas.expand(-1, -1, hidden_states.shape[2])
+        if self.training:
+            self._auxiliary_loss = self.compute_loss(alphas)
+            value_states = hidden_states.abs()  # ensure value_states is positive
+            normalized_value_states, scale_factor, shift_factor = self.normalizer(value_states)
+
+            betas = torch.reciprocal(normalized_value_states)
+            sign_states = torch.sign(hidden_states)
+            gamma_dist = torch.distributions.gamma.Gamma(shaped_alphas, betas)
+            samples = gamma_dist.rsample()
+            # Restore the original scale
+            if shift_factor is not None:
+                time_states = samples * sign_states / scale_factor - shift_factor / scale_factor
+            else:
+                time_states = samples * sign_states / scale_factor
+        else:
+            time_states = shaped_alphas * hidden_states # directly use the expectation
+        
+        if self.return_attn:
+            # print(alphas.shape)
+            if time_states.shape[-2] == 1:
+                pass # seqlen == 1 indicates it is caching
+            else:
+                self._attn = torch.exp(-time_states).mean(dim=-1).detach().cpu()
+
+        return time_states
+
+class BernoulliIB(nn.Module):
+    def __init__(self, hidden_size, temp=1, thetas=None, max_seqlen=1024, return_attn=False, **kwargs) -> None:
+        super().__init__()
+        self.epoch_frac = 0
+        self.epoch_threshold = 0
+        self.temp = temp
+        self.thetas = thetas
+        self.hidden_size = hidden_size
+        # self.attributor = nn.Linear(hidden_size, 1, bias=True)
+        self._auxiliary_loss = 0
+        self.max_seqlen = 4096
+        self.return_attn = return_attn
+        self._attn = None
+
+    def init_thetas(self, attn):
+        # Create a tensor with sequence positions [0, 1, ..., length-1]
+        # length = attn.shape[1]
+        seq_len = attn.shape[1]
+        if seq_len <= self.max_seqlen:
+            positions = torch.arange(self.max_seqlen).float().to(attn.device)
+        else:
+            positions = torch.arange(self.max_seqlen - seq_len, self.max_seqlen).float().to(attn.device)
+        
+        # Define the exponential decay function
+        # decay_factor = 0.3 + 0.4 * torch.exp(positions / length - 1) # extrapolable distance-decay (-\infty: 0.3, 0: 0.5, length: 0.7)
+        decay_factor = 0.8 - 0.6 * torch.exp(positions / self.max_seqlen - 1) # extrapolable distance-balance (-\infty: 0.7, 0: 0.5, length: 0.3)
+        
+        # Make the decay factor repeat across the batch dimension
+        self.thetas = decay_factor
+        return decay_factor
+
+    # def get_token_saliency(self):
+    #     attn = self._attn
+    #     self._attn = None
+    #     return attn
+
+    def get_auxiliary_loss(self):
+        loss = self._auxiliary_loss
+        self._auxiliary_loss = 0.0
+        return loss
+
+    def compute_loss(self, att, epsilon=1e-6):
+        if self.thetas is None:
+            thetas = self.init_thetas(att)  # length is the second dimension of att
+            print(f"Bernoulli prior theta first: {self.thetas[0]:.2f}, last: {self.thetas[-1]:.2f}, size: {self.thetas.size(0)}")
+
+        if self.thetas.size(0) >= att.size(1):
+            thetas = self.thetas[-att.size(1):]
+        elif self.thetas.size(0) < att.size(1):
+            thetas = self.init_thetas(att)
+            print(f"Bernoulli prior theta first: {self.thetas[0]:.2f}, last: {self.thetas[-1]:.2f}, size: {self.thetas.size(0)}")
+        
+        thetas = thetas.unsqueeze(-1).expand(att.shape)
+        # Calculate the regularization loss
+        reg_loss = (att * torch.log(att / thetas + epsilon) +
+                    (1 - att) * torch.log((1 - att) / (1 - thetas + epsilon) + epsilon)).mean()
+        
+        return reg_loss
+
+    def forward(self, states):
+        # hidden_states shape [bsz, seq_length, dimension]
+        hidden_states, attn = torch.split(
+                    states,
+                    [self.hidden_size, 1],
+                    dim=-1
+                )
+        # attn = self.attributor(hidden_states)
+        #! HACK: for fintuning, we zero-init and re-paramerize as plus 1
+        attn = attn + torch.tensor(1.0)
+        if self.epoch_frac < self.epoch_threshold:
+            return hidden_states
+        if self.training:
+            # gumble soft-max
+            random_noise = torch.empty_like(attn).uniform_(1e-10, 1 - 1e-10)
+            random_noise = torch.log(random_noise) - torch.log(1.0 - random_noise)
+            attn_bern = ((attn + random_noise) / self.temp).sigmoid()
+        else:
+            attn_bern = (attn).sigmoid()
+        self._auxiliary_loss = self.compute_loss(attn_bern)
+        if self.return_attn:
+            if attn_bern.shape[-2] == 1:
+                pass # seqlen == 1 indicates it is caching
+            else:
+                self._attn = attn_bern.detach().cpu()
+        return hidden_states * attn_bern
+
+
+class Mamba2Mixer(nn.Module):
+    """
+    Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
+    A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
+    ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
+    and is why Mamba is called **selective** state spaces)
+    """
+
+    def __init__(self, config: IBS2Config, layer_idx: int):
+        super().__init__()
+        self.num_heads = config.num_heads
+        self.hidden_size = config.hidden_size
+        self.ssm_state_size = config.state_size
+        self.conv_kernel_size = config.conv_kernel
+        self.intermediate_size = int(config.expand * self.hidden_size)
+        self.time_step_rank = int(config.time_step_rank)
+        self.layer_idx = layer_idx
+        self.use_conv_bias = config.use_conv_bias
+        self.activation = config.hidden_act
+        self.act = ACT2FN[config.hidden_act]
+
+        self.layer_norm_epsilon = config.layer_norm_epsilon
+        self.rms_norm = config.rms_norm
+
+        self.n_groups = config.n_groups
+        self.head_dim = config.head_dim
+        self.chunk_size = config.chunk_size
+
+        self.time_step_limit = config.time_step_limit
+        self.time_step_min = config.time_step_min
+        self.time_step_max = config.time_step_max
+
+        self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            bias=config.use_conv_bias,
+            kernel_size=config.conv_kernel,
+            groups=self.conv_dim,
+            padding=config.conv_kernel - 1,
+        )
+
+        # projection of the input hidden states
+        #! HACK ib_dim
+        self._attn = None
+        self.return_attn = config.return_attn
+        assert config.ib_type in ['bernoulli', 'gamma'], "Invalid IB Prior."
+        IB_cls = BernoulliIB if config.ib_type == 'bernoulli' else GammaIB if config.ib_type == 'gamma' else None
+        self.ib4dt = IB_cls(self.num_heads, return_attn=config.return_attn) if self.layer_idx in [0, 31, 63] else None
+
+        self.ib_proj = nn.Linear(
+            self.hidden_size,
+            1,
+            bias=False,
+        ) if self.ib4dt else None
+        projection_size = self.intermediate_size + self.conv_dim + self.num_heads
+        self.in_proj = nn.Linear(
+            self.hidden_size,
+            projection_size,
+            bias=config.use_bias,
+        )
+        # selective projection used to make dt, B and C input dependant
+
+        # time step projection (discretization)
+        # instantiate once and copy inv_dt in init_weights of PretrainedModel
+        self.dt_bias = nn.Parameter(torch.ones(self.num_heads))
+
+        # S4D real initialization. These are not discretized!
+        # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded
+        A = torch.arange(1, self.num_heads + 1)
+        self.A_log = nn.Parameter(torch.log(A))
+        self.A_log._no_weight_decay = True
+        self.norm = MambaRMSNormGated(self.intermediate_size, eps=self.layer_norm_epsilon)
+        self.D = nn.Parameter(torch.ones(self.num_heads))
+        self.D._no_weight_decay = True
+
+        self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
+        self.use_bias = config.use_bias
+
+        if not is_fast_path_available:
+            logger.warning_once(
+                "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+                " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
+                " https://github.com/Dao-AILab/causal-conv1d"
+            )
+
+    def get_token_saliency(self):
+        attn = self._attn
+        self._attn = None
+        return attn
+
+    def cuda_kernels_forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_params: Optional[Mamba2Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ):
+        # 1. Gated MLP's linear projection
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+        projected_states = self.in_proj(hidden_states)
+
+        #! HACK IBS apply
+        if self.ib4dt:
+            ib_state = self.ib_proj(hidden_states)
+            dim = self.head_dim * self.num_heads 
+            zx, BC, dt = torch.split(projected_states, [dim * 2, + self.n_groups * self.ssm_state_size * 2, self.num_heads], dim=-1)
+            dt = self.ib4dt(torch.cat([dt, ib_state], dim=-1))
+            projected_states = torch.cat([zx, BC, dt], dim=-1)
+            if self.return_attn and dt.shape[-2] != 1:
+                dt_plus = nn.functional.softplus(dt + self.dt_bias)
+                dA = (dt_plus * (-torch.exp(self.A_log.float())))
+                # dA = torch.exp(dA)
+                attn = dA.mean(dim=-1) # - 0.1 * dA.std(dim=-1) # attn shape [batch_size, seqlen]
+                self._attn = attn
+
+        # Set up dimensions for reshapes later
+        batch_size, seq_len, _ = hidden_states.shape
+        groups_time_state_size = self.n_groups * self.ssm_state_size
+        d_mlp = (
+            projected_states.shape[-1]
+            - 2 * self.intermediate_size
+            - 2 * self.n_groups * self.ssm_state_size
+            - self.num_heads
+        ) // 2
+
+        # Single step calculations via cache
+        if cache_params is not None and cache_position is not None and cache_position[0] > 0:
+            _, _, gate, hidden_states_B_C, dt = projected_states.squeeze(1).split(
+                [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
+            )
+
+            # 2. Convolution sequence transformation
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                cache_params.conv_states[self.layer_idx],
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                self.activation,
+            )
+
+            hidden_states, B, C = torch.split(
+                hidden_states_B_C,
+                [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+                dim=-1,
+            )
+
+            # 3. SSM transformation
+            A = -torch.exp(self.A_log.float())  # (nheads,)
+            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            dt = dt[:, :, None].expand(-1, -1, self.head_dim)
+            dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
+            D = self.D[:, None, ...].expand(-1, self.head_dim)
+            B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups)
+            C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups)
+            hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
+            hidden_states = selective_state_update(
+                cache_params.ssm_states[self.layer_idx],
+                hidden_states_reshaped,
+                dt,
+                A,
+                B,
+                C,
+                D,
+                z=None,
+                dt_bias=dt_bias,
+                dt_softplus=True,
+            )
+            hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim)
+            hidden_states = self.norm(hidden_states, gate)
+
+            # 4. Final linear projection
+            out = self.out_proj(hidden_states)[:, None, ...]
+
+        # Fused calculations or step by step if no initialized cache is found
+        else:
+            A = -torch.exp(self.A_log.float())  # (num_heads) or (intermediate_size, state_size)
+            dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
+
+            # 2-4. Fused kernel for conv1d, SSM, and the final projection
+            if self.training and cache_params is None:
+                out = mamba_split_conv1d_scan_combined(
+                    projected_states,
+                    self.conv1d.weight.squeeze(1),
+                    self.conv1d.bias,
+                    self.dt_bias,
+                    A,
+                    D=self.D,
+                    chunk_size=self.chunk_size,
+                    seq_idx=None,  # was seq_idx
+                    activation=self.activation,
+                    rmsnorm_weight=self.norm.weight,
+                    rmsnorm_eps=self.norm.variance_epsilon,
+                    outproj_weight=self.out_proj.weight,
+                    outproj_bias=self.out_proj.bias,
+                    headdim=self.head_dim,
+                    ngroups=self.n_groups,
+                    norm_before_gate=False,
+                    return_final_states=False,
+                    **dt_limit_kwargs,
+                )
+
+            else:
+                _, _, gate, hidden_states_B_C, dt = projected_states.split(
+                    [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
+                )
+
+                # 2. Convolution sequence transformation
+                # Init cache
+                if cache_params is not None:
+                    hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2)
+                    conv_states = nn.functional.pad(
+                        hidden_states_B_C_transposed,
+                        (cache_params.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0),
+                    )
+                    cache_params.update_conv_state(
+                        layer_idx=self.layer_idx, new_conv_state=conv_states, cache_init=True
+                    )
+
+                if self.activation not in ["silu", "swish"]:
+                    hidden_states_B_C = self.act(
+                        self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2)
+                    )
+                else:
+                    hidden_states_B_C = causal_conv1d_fn(
+                        x=hidden_states_B_C.transpose(1, 2),
+                        weight=self.conv1d.weight.squeeze(1),
+                        bias=self.conv1d.bias,
+                        activation=self.activation,
+                    ).transpose(1, 2)
+
+                hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+                hidden_states, B, C = torch.split(
+                    hidden_states_B_C,
+                    [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+                    dim=-1,
+                )
+
+                # 3. SSM transformation
+                scan_output, ssm_state = mamba_chunk_scan_combined(
+                    hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+                    dt,
+                    A,
+                    B.view(batch_size, seq_len, self.n_groups, -1),
+                    C.view(batch_size, seq_len, self.n_groups, -1),
+                    chunk_size=self.chunk_size,
+                    D=self.D,
+                    z=None,
+                    seq_idx=None,
+                    return_final_states=True,
+                    dt_bias=self.dt_bias,
+                    dt_softplus=True,
+                    **dt_limit_kwargs,
+                )
+
+                # Init cache
+                if ssm_state is not None and cache_params is not None:
+                    cache_params.update_ssm_state(layer_idx=self.layer_idx, new_ssm_state=ssm_state)
+
+                scan_output = scan_output.view(batch_size, seq_len, -1)
+                # Multiply "gate" branch and apply extra normalization layer
+                scan_output = self.norm(scan_output, gate)
+
+                # 4. Final linear projection
+                out = self.out_proj(scan_output)
+        return out
+
+    # fmt: off
+    def torch_forward(self, input_states, cache_params: Optional[Mamba2Cache]=None, cache_position:Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None):
+        batch_size, seq_len, _ = input_states.shape
+        dtype = input_states.dtype
+
+        # 1. Gated MLP's linear projection
+        input_states = apply_mask_to_padding_states(input_states, attention_mask)
+        projected_states = self.in_proj(input_states)
+        
+        if self.ib4dt:
+            attn = self.ib_proj(input_states)
+            dim = self.head_dim * self.num_heads 
+            zx, BC, dt = torch.split(projected_states, [dim * 2, + self.n_groups * self.ssm_state_size * 2, self.num_heads], dim=-1)
+            dt = self.ib4dt(torch.cat([dt, attn], dim=-1))
+            projected_states = torch.cat([zx, BC, dt], dim=-1)
+
+        d_mlp = (projected_states.shape[-1] - 2 * self.intermediate_size - 2 * self.n_groups * self.ssm_state_size-self.num_heads) // 2
+        _, _, gate, hidden_states_B_C, dt = projected_states.split(
+                [d_mlp, d_mlp, self.intermediate_size,  self.conv_dim, self.num_heads], dim=-1
+        )
+
+        # 2. Convolution sequence transformation
+        if cache_params is not None and cache_position is not None and cache_position[0] > 0:
+            cache_params.update_conv_state(layer_idx=self.layer_idx, new_conv_state=hidden_states_B_C, cache_init=False)
+
+            # We need to guarantee that anything regarding the cache is on the same device
+            conv_states = cache_params.conv_states[self.layer_idx].to(device=self.conv1d.weight.device)
+
+            hidden_states_B_C = torch.sum(
+                conv_states * self.conv1d.weight.squeeze(1), dim=-1
+            )
+            if self.use_conv_bias:
+                hidden_states_B_C = hidden_states_B_C + self.conv1d.bias
+            hidden_states_B_C = self.act(hidden_states_B_C)
+        else:
+            # Init cache
+            if cache_params is not None:
+                hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2)
+                conv_states = nn.functional.pad(
+                    hidden_states_B_C_transposed, (cache_params.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0)
+                )
+                cache_params.update_conv_state(layer_idx=self.layer_idx, new_conv_state=conv_states, cache_init=True)
+
+            hidden_states_B_C = self.act(self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2))
+
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+        hidden_states, B, C = torch.split(
+            hidden_states_B_C,
+            [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
+            dim=-1
+        )
+
+        # 3. SSM transformation
+        A = -torch.exp(self.A_log.float())                            # [num_heads]
+        if cache_params is not None and cache_position is not None and cache_position[0] > 0:
+            # We need to guarantee that anything regarding the cache is on the same device
+            cache_device = cache_params.ssm_states.device
+
+            # Note: there is no need to pad parameter matrices here, as there is just one new token
+            # for batched generation
+            dt = dt[:, 0, :][:, None, ...]
+            dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim)
+            # [num_heads] -> [num_heads, head_dim]
+            dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim)
+
+            dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))
+            dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])
+            A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            # [bsz, num_heads, head_dim, state_size]
+            dA = (torch.exp(dt[..., None] * A)).to(device=cache_device)
+
+            # Discretize B
+            # [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] ->
+            # -> [bsz, n_groups, group to head repetition factor, state_size] -> [bsz, num_heads, state_size]
+            B = B.reshape(batch_size, self.n_groups, -1)[..., None, :]
+            B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
+            B = B.reshape(batch_size, -1, B.shape[-1])
+            # [bsz, num_heads, head_dim, state_size]
+            dB = dt[..., None] * B[..., None, :]
+
+            # Discretize x into dB
+            # [bsz, intermediate_size] -> [bsz, num_heads, head_dim]
+            hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
+            dBx = (dB * hidden_states[..., None]).to(device=cache_device)
+
+            # State calculation
+            cache_params.update_ssm_state(
+                layer_idx=self.layer_idx,
+                new_ssm_state=cache_params.ssm_states[self.layer_idx] * dA + dBx
+            )
+
+            # Subsequent output
+            # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size]
+            C = C.reshape(batch_size, self.n_groups, -1)[..., None, :]
+            C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous()
+            C = C.reshape(batch_size, -1, C.shape[-1])
+            # [bsz, num_heads, head_dim]
+
+            ssm_states = cache_params.ssm_states[self.layer_idx].to(device=C.device, dtype=C.dtype)  # Shape: [b, h, d, n]
+            # Reshape ssm_states to merge the first two dimensions
+            ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size)  # Shape: [b*h, d, n]
+            C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1)  # Shape: [b*h, n, 1]
+            y = torch.bmm(ssm_states_reshaped, C_reshaped)
+            y = y.view(batch_size, self.num_heads, self.head_dim)
+
+            # D skip connection
+            # [num_heads] -> [num_heads, head_dim]
+            D = self.D[..., None].expand(self.D.shape[0], self.head_dim)
+            y = (y + hidden_states * D).to(y.dtype)
+
+            # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size]
+            y = y.reshape(batch_size, -1)[:, None, ...]
+        else:
+            # begin ssd naive implementation without einsums
+            dt = nn.functional.softplus(dt + self.dt_bias)
+            dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])
+            hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float()
+            B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
+            C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
+            B = B.repeat(1, 1, self.num_heads // self.n_groups, 1)
+            C = C.repeat(1, 1, self.num_heads // self.n_groups, 1)
+            pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size
+
+            D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+            # Discretize x and A
+            hidden_states = hidden_states * dt[..., None]
+            A = A.to(hidden_states.dtype) * dt
+
+            # Rearrange into blocks/chunks
+            hidden_states, A, B, C = [reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C)]
+
+            # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size]
+            A = A.permute(0, 3, 1, 2)
+            A_cumsum = torch.cumsum(A, dim=-1)
+
+            # 1. Compute the output for each intra-chunk (diagonal blocks)
+            # This is the analog of a causal mask
+            L = torch.exp(segment_sum(A))
+
+            # Contraction of C and B to get G (attention-weights like)
+            G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]  # shape: (b, c, l, s, h, n)
+            G = G_intermediate.sum(dim=-1)  # shape: (b, c, l, s, h)
+
+            # Compute M, equivalent to applying attention mask to weights
+            M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]
+            M = M_intermediate.sum(dim=-1)
+
+            # Compute Y_diag (apply to values)
+            Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3)
+
+            # 2. Compute the state for each intra-chunk
+            # (right term of low-rank factorization of off-diagonal blocks; B terms)
+            decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
+            B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None]
+            states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2)
+
+            # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
+            # (middle term of factorization of off-diag blocks; A terms)
+            if cache_params is not None and cache_position is not None and cache_position[0] > 0:
+                previous_states = cache_params.ssm_states[self.layer_idx][:, None, ...].to(device=states.device)
+            else:
+                previous_states = torch.zeros_like(states[:, :1])
+            states = torch.cat([previous_states, states], dim=1)
+            decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0))))
+            decay_chunk = decay_chunk.transpose(1, 3)
+            new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+            states, ssm_state = new_states[:, :-1], new_states[:, -1]
+
+            # 4. Compute state -> output conversion per chunk
+            # (left term of low-rank factorization of off-diagonal blocks; C terms)
+            state_decay_out = torch.exp(A_cumsum)
+            C_times_states = (C[..., None, :] * states[:, :, None, ...])
+            state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1)
+            Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None])
+
+            # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+            y = Y_diag + Y_off
+            # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim]
+            y = y.reshape(batch_size, -1, self.num_heads, self.head_dim)
+
+            y = y + D_residual
+            # Cutting off padded chunks
+            if pad_size > 0:
+                y = y[:, :seq_len, :, :]
+            y = y.reshape(batch_size, seq_len, -1)
+
+            # Init cache
+            if ssm_state is not None and cache_params is not None:
+                cache_params.update_ssm_state(layer_idx=self.layer_idx, new_ssm_state=ssm_state)
+
+        scan_output = self.norm(y, gate)
+
+        # end ssd naive
+
+        # 4. Final linear projection
+        contextualized_states = self.out_proj(scan_output.to(dtype))  # [batch, seq_len, hidden_size]
+        return contextualized_states
+    # fmt: on
+
+    def forward(
+        self,
+        hidden_states,
+        cache_params: Optional[Mamba2Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ):
+        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type:
+            return self.cuda_kernels_forward(hidden_states, cache_params, cache_position, attention_mask)
+        dtype = hidden_states.dtype
+        if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+            # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
+            hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+        return self.torch_forward(hidden_states, cache_params, cache_position, attention_mask)
+
+
+class Mamba2RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        Mamba2RMSNorm is equivalent to T5LayerNorm and LlamaRMSNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class IBS2Block(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.residual_in_fp32 = config.residual_in_fp32
+        self.norm = Mamba2RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+        self.mixer = Mamba2Mixer(config, layer_idx=layer_idx)
+
+    def forward(
+        self,
+        hidden_states,
+        cache_params: Optional[Mamba2Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ):
+        residual = hidden_states
+        hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+        if self.residual_in_fp32:
+            residual = residual.to(torch.float32)
+
+        hidden_states = self.mixer(
+            hidden_states, cache_params=cache_params, cache_position=cache_position, attention_mask=attention_mask
+        )
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class Mamba2PreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = IBS2Config
+    base_model_prefix = "backbone"
+    _no_split_modules = ["Mamba2Block"]
+    supports_gradient_checkpointing = True
+    _is_stateful = True
+
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        if isinstance(module, Mamba2Mixer):
+            #! HACK
+            if getattr(module, "ib_proj"):
+                nn.init.zeros_(module.ib_proj.weight)
+
+            module.A_log._no_weight_decay = True
+            module.D._no_weight_decay = True
+
+            dt = torch.exp(
+                torch.rand(self.config.num_heads)
+                * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min))
+                + math.log(self.config.time_step_min)
+            ).clamp(min=self.config.time_step_floor)
+
+            # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+            inv_dt = dt + torch.log(-torch.expm1(-dt))
+            with torch.no_grad():
+                module.dt_bias.copy_(inv_dt)
+            module.dt_bias._no_reinit = True
+
+        if isinstance(module, nn.Linear):
+            if module.bias is not None:
+                if not getattr(module.bias, "_no_reinit", False):
+                    nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=self.config.initializer_range)
+
+        if self.config.rescale_prenorm_residual:
+            # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+            #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+            #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+            #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+            #
+            # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+            for name, p in module.named_parameters():
+                if name in ["out_proj.weight"]:
+                    # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                    # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                    # We need to reinit p since this code could be called multiple times
+                    # Having just p *= scale would repeatedly scale it down
+                    nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                    with torch.no_grad():
+                        p /= math.sqrt(self.config.num_hidden_layers)
+
+
+@dataclass
+# Copied from transformers.models.mamba.modeling_mamba.MambaOutput with MAMBA->MAMBA2,Mamba->Mamba2
+class Mamba2Output(ModelOutput):
+    """
+    Class for the MAMBA2 model outputs.
+
+    Args:
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        cache_params (`Mamba2Cache`):
+            The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
+            avoid providing the old `input_ids`.
+
+            Includes both the State space model state matrices after the selective scan, and the Convolutional states
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+    """
+
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    cache_params: Optional[Mamba2Cache] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+# Copied from transformers.models.mamba.modeling_mamba.MambaCausalLMOutput with Mamba->Mamba2
+class Mamba2CausalLMOutput(ModelOutput):
+    """
+    Base class for causal language model (or autoregressive) outputs.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Language modeling loss (for next-token prediction).
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        cache_params (`Mamba2Cache`):
+            The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
+            avoid providing the old `input_ids`.
+
+            Includes both the State space model state matrices after the selective scan, and the Convolutional states
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    cache_params: Optional[Mamba2Cache] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+
+
+MAMBA2_START_DOCSTRING = r"""
+
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`Mamba2Config`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+MAMBA2_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+            Indices of input sequence tokens in the vocabulary.
+
+            If `cache_params.seqlen_offset>0`, only `input_ids` that do not have their past calculated should be passed as
+            `input_ids`.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        cache_params (`Mamba2Cache`, *optional*):
+            If passed along, the model uses the previous state in all the blocks (which will give the output for the
+            `input_ids` provided as if the model add `state_input_ids + input_ids` as context).
+        use_cache (`bool`, *optional*):
+            If set to `True`, the `cache_params` is returned and can be used to quickly generate the next logits.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        cache_position (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            The position of the current input in the cache. This is used to ensure that the cache is correctly updated.
+            If `cache_params` is passed, `cache_position` should also be passed.
+        attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+"""
+
+
+@add_start_docstrings(
+    "The bare MAMBA2 Model transformer outputting raw hidden-states without any specific head on top.",
+    MAMBA2_START_DOCSTRING,
+)
+class IBS2Model(Mamba2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([IBS2Block(config, layer_idx=idx) for idx in range(config.num_hidden_layers)])
+
+        self.gradient_checkpointing = False
+        self.norm_f = Mamba2RMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+        # Initialize weights and apply final processing
+        self._register_load_state_dict_pre_hook(self.load_hook)
+        self.post_init()
+
+    def load_hook(self, state_dict, prefix, *args):
+        for k in state_dict:
+            if "embedding." in k:
+                state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k)
+                break
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, new_embeddings):
+        self.embeddings = new_embeddings
+
+    @add_start_docstrings_to_model_forward(MAMBA2_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=Mamba2Output,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.LongTensor] = None,
+        cache_params: Optional[Mamba2Cache] = None,
+        use_cache: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, Mamba2Output]:
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):  # ^ is python for xor
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embeddings(input_ids)
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            use_cache = False
+
+        if use_cache:
+            if cache_params is None:
+                cache_params = Mamba2Cache(
+                    self.config, inputs_embeds.size(0), device=inputs_embeds.device, dtype=inputs_embeds.dtype
+                )
+                cache_position = torch.arange(0, self.config.conv_kernel, device=inputs_embeds.device)
+            elif cache_position is None:
+                # cases when we do manual forward instead of using `model.generate` which will initiate
+                # `cache_position` and makes sure it is not None, throw error here instead of doing some
+                # hack to conjecture the current cache position
+                raise ValueError(
+                    "You have to specify the `cache_position` manually when `use_cache=True` and `cache_params` is passed, "
+                    "you don't have to pass a `cache_params` if you are in prefilling stage because in that case it will "
+                    "be initialized for you automatically"
+                )
+        else:
+            cache_params = None
+
+        hidden_states = inputs_embeds
+        all_hidden_states = () if output_hidden_states else None
+        for mixer_block in self.layers:
+            if self.gradient_checkpointing and self.training:
+                hidden_states = self._gradient_checkpointing_func(
+                    mixer_block.__call__, hidden_states, cache_params, cache_position, attention_mask
+                )
+            else:
+                hidden_states = mixer_block(
+                    hidden_states,
+                    cache_params=cache_params,
+                    cache_position=cache_position,
+                    attention_mask=attention_mask,
+                )
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+        hidden_states = self.norm_f(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, cache_params, all_hidden_states] if v is not None)
+
+        return Mamba2Output(
+            last_hidden_state=hidden_states,
+            cache_params=cache_params if use_cache else None,
+            hidden_states=all_hidden_states,
+        )
+
+class IBS2ForClassification(Mamba2PreTrainedModel):
+    _tied_weights_keys = []
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.backbone = IBS2Model(config)
+        self.cls_head = nn.Linear(config.hidden_size, config.num_classes, bias=False)
+        # Initialize weights and apply final processing
+        self.post_init()
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_params: Optional[Mamba2Cache] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,  # for now we need this for generation
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        mamba2_outputs = self.backbone(
+            input_ids,
+            cache_params=cache_params,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            attention_mask=attention_mask,
+        )
+        hidden_states = mamba2_outputs[0]
+
+        logits = self.cls_head(hidden_states.to(self.cls_head.weight.dtype)).float()
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + mamba2_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return loss
+        # return Mamba2CausalLMOutput(
+        #     loss=loss,
+        #     logits=logits,
+        #     cache_params=mamba2_outputs.cache_params,
+        #     hidden_states=mamba2_outputs.hidden_states,
+        # )
+
+@add_start_docstrings(
+    """
+    The MAMBA2 Model transformer with a language modeling head on top (linear layer with weights not tied to the input
+    embeddings).
+    """,
+    MAMBA2_START_DOCSTRING,
+)
+class IBS2ForCausalLM(Mamba2PreTrainedModel, GenerationMixin):
+    _tied_weights_keys = []
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.backbone = IBS2Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def get_input_embeddings(self):
+        return self.backbone.get_input_embeddings()
+
+    def set_input_embeddings(self, new_embeddings):
+        return self.backbone.set_input_embeddings(new_embeddings)
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        inputs_embeds=None,
+        use_cache=None,
+        cache_params: Optional[Mamba2Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        # Overwitten -- uses `cache_params` as opposed to `past_key_values`
+
+        if use_cache:
+            # `cache_position` should have been initialized in `generate`
+            if cache_position is None:
+                raise ValueError(
+                    "`cache_position` should not be None as it should have been initialized in "
+                    "`model.generate`, you are responsible for passing in a valid `cache_position` if "
+                    "you are calling `prepare_inputs_for_generation` directly with `use_cache=True`"
+                )
+            if cache_position[0] > 0:
+                input_ids = input_ids[:, -1][..., None]
+
+                if attention_mask is not None:
+                    attention_mask = None
+            else:
+                # we initialize the `cache_position` to full size of `conv_states` at prefill stage
+                # considering padding will be applied when input length is shorter, and truncation
+                # will be applied when it is longer, so it will be equivalent to always have it match
+                # the length of `cache_params.conv_states`, which is `config.conv_kernel`
+                cache_position = torch.arange(0, self.config.conv_kernel, device=input_ids.device)
+
+        if inputs_embeds is not None and cache_params is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "attention_mask": attention_mask,
+                "cache_params": cache_params,
+                "use_cache": use_cache,
+                "cache_position": cache_position,
+            }
+        )
+        return model_inputs
+
+    @add_start_docstrings_to_model_forward(MAMBA2_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=Mamba2CausalLMOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_params: Optional[Mamba2Cache] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,  # for now we need this for generation
+    ) -> Union[Tuple, Mamba2CausalLMOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        mamba2_outputs = self.backbone(
+            input_ids,
+            cache_params=cache_params,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            attention_mask=attention_mask,
+        )
+        hidden_states = mamba2_outputs[0]
+
+        logits = self.lm_head(hidden_states.to(self.lm_head.weight.dtype)).float()
+
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + mamba2_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return Mamba2CausalLMOutput(
+            loss=loss,
+            logits=logits,
+            cache_params=mamba2_outputs.cache_params,
+            hidden_states=mamba2_outputs.hidden_states,
+        )
+
+
+__all__ = ["IBS2ForCausalLM", "IBS2Model", "Mamba2PreTrainedModel", "IBS2Block", "IBS2ForClassification"]
+

pytorch_model-00001-of-00006.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f87d5313fd2baea17acf628de1b2922c42563b4999c10e12dc3f8bb9073ee75
+size 4922747462

pytorch_model-00002-of-00006.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:220fee199c4c00bd513b86585915fa89057a080aa9af6a0a6742a7b48711813a
+size 4824198060

pytorch_model-00003-of-00006.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:245bf3b887b6ce842a330c220e26d5c5d74452a9a0feaa2dec460de4dea24242
+size 4824214794

pytorch_model-00004-of-00006.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:834b44b5afc90469d12f4072c558db13aec273d6659a0fec88cd04d730917614
+size 4824198060

pytorch_model-00005-of-00006.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6bf50bc90604dc0f7523a05dc68130e1e82ac89b9560208cccc5aca7cdc7814f
+size 4824198060

pytorch_model-00006-of-00006.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a981ad272e9ed171dd03a82b720ade1b1d35724dc0b61867c7a778fedc62f57
+size 4922314420

pytorch_model.bin.index.json

@@ -0,0 +1,589 @@
+{
+  "metadata": {
+    "total_size": 29141663744
+  },
+  "weight_map": {
+    "backbone.embeddings.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.0.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.0.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.0.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.0.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.0.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.0.mixer.ib_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.0.mixer.in_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.0.mixer.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.0.mixer.out_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.0.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.1.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.1.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.1.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.1.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.1.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.1.mixer.in_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.1.mixer.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.1.mixer.out_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.1.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.10.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.10.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.10.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.10.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.10.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.10.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.10.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.10.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.10.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.11.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.11.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.11.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.11.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.11.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.11.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.11.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.11.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.11.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.12.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.12.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.12.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.12.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.12.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.12.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.12.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.12.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.12.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.13.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.13.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.13.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.13.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.13.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.13.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.13.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.13.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.13.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.14.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.14.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.14.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.14.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.14.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.14.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.14.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.14.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.14.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.15.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.15.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.15.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.15.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.15.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.15.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.15.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.15.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.15.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.16.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.16.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.16.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.16.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.16.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.16.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.16.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.16.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.16.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.17.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.17.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.17.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.17.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.17.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.17.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.17.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.17.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.17.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.18.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.18.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.18.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.18.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.18.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.18.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.18.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.18.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.18.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.19.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.19.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.19.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.19.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.19.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.19.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.19.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.19.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.19.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.2.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.2.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.2.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.2.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.2.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.2.mixer.in_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.2.mixer.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.2.mixer.out_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.2.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.20.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.20.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.20.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.20.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.20.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.20.mixer.in_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.20.mixer.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.20.mixer.out_proj.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.20.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.21.mixer.A_log": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.21.mixer.D": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.21.mixer.conv1d.bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.21.mixer.conv1d.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.21.mixer.dt_bias": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.21.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.21.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.21.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.21.norm.weight": "pytorch_model-00002-of-00006.bin",
+    "backbone.layers.22.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.22.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.22.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.22.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.22.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.22.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.22.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.22.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.22.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.23.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.23.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.23.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.23.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.23.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.23.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.23.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.23.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.23.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.24.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.24.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.24.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.24.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.24.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.24.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.24.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.24.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.24.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.25.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.25.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.25.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.25.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.25.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.25.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.25.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.25.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.25.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.26.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.26.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.26.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.26.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.26.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.26.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.26.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.26.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.26.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.27.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.27.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.27.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.27.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.27.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.27.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.27.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.27.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.27.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.28.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.28.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.28.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.28.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.28.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.28.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.28.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.28.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.28.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.29.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.29.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.29.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.29.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.29.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.29.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.29.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.29.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.29.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.3.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.3.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.3.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.3.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.3.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.3.mixer.in_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.3.mixer.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.3.mixer.out_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.3.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.30.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.30.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.30.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.30.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.30.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.30.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.30.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.30.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.30.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.31.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.31.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.31.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.31.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.31.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.31.mixer.ib_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.31.mixer.in_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.31.mixer.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.31.mixer.out_proj.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.31.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.32.mixer.A_log": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.32.mixer.D": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.32.mixer.conv1d.bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.32.mixer.conv1d.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.32.mixer.dt_bias": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.32.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.32.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.32.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.32.norm.weight": "pytorch_model-00003-of-00006.bin",
+    "backbone.layers.33.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.33.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.33.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.33.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.33.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.33.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.33.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.33.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.33.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.34.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.34.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.34.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.34.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.34.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.34.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.34.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.34.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.34.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.35.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.35.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.35.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.35.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.35.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.35.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.35.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.35.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.35.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.36.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.36.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.36.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.36.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.36.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.36.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.36.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.36.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.36.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.37.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.37.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.37.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.37.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.37.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.37.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.37.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.37.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.37.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.38.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.38.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.38.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.38.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.38.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.38.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.38.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.38.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.38.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.39.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.39.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.39.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.39.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.39.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.39.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.39.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.39.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.39.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.4.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.4.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.4.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.4.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.4.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.4.mixer.in_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.4.mixer.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.4.mixer.out_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.4.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.40.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.40.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.40.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.40.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.40.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.40.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.40.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.40.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.40.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.41.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.41.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.41.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.41.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.41.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.41.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.41.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.41.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.41.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.42.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.42.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.42.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.42.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.42.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.42.mixer.in_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.42.mixer.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.42.mixer.out_proj.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.42.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.43.mixer.A_log": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.43.mixer.D": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.43.mixer.conv1d.bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.43.mixer.conv1d.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.43.mixer.dt_bias": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.43.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.43.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.43.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.43.norm.weight": "pytorch_model-00004-of-00006.bin",
+    "backbone.layers.44.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.44.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.44.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.44.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.44.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.44.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.44.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.44.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.44.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.45.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.45.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.45.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.45.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.45.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.45.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.45.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.45.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.45.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.46.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.46.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.46.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.46.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.46.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.46.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.46.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.46.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.46.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.47.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.47.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.47.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.47.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.47.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.47.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.47.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.47.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.47.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.48.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.48.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.48.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.48.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.48.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.48.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.48.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.48.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.48.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.49.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.49.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.49.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.49.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.49.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.49.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.49.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.49.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.49.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.5.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.5.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.5.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.5.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.5.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.5.mixer.in_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.5.mixer.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.5.mixer.out_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.5.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.50.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.50.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.50.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.50.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.50.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.50.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.50.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.50.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.50.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.51.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.51.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.51.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.51.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.51.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.51.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.51.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.51.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.51.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.52.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.52.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.52.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.52.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.52.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.52.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.52.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.52.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.52.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.53.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.53.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.53.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.53.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.53.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.53.mixer.in_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.53.mixer.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.53.mixer.out_proj.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.53.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.54.mixer.A_log": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.54.mixer.D": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.54.mixer.conv1d.bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.54.mixer.conv1d.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.54.mixer.dt_bias": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.54.mixer.in_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.54.mixer.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.54.mixer.out_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.54.norm.weight": "pytorch_model-00005-of-00006.bin",
+    "backbone.layers.55.mixer.A_log": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.55.mixer.D": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.55.mixer.conv1d.bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.55.mixer.conv1d.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.55.mixer.dt_bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.55.mixer.in_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.55.mixer.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.55.mixer.out_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.55.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.56.mixer.A_log": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.56.mixer.D": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.56.mixer.conv1d.bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.56.mixer.conv1d.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.56.mixer.dt_bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.56.mixer.in_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.56.mixer.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.56.mixer.out_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.56.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.57.mixer.A_log": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.57.mixer.D": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.57.mixer.conv1d.bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.57.mixer.conv1d.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.57.mixer.dt_bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.57.mixer.in_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.57.mixer.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.57.mixer.out_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.57.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.58.mixer.A_log": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.58.mixer.D": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.58.mixer.conv1d.bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.58.mixer.conv1d.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.58.mixer.dt_bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.58.mixer.in_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.58.mixer.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.58.mixer.out_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.58.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.59.mixer.A_log": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.59.mixer.D": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.59.mixer.conv1d.bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.59.mixer.conv1d.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.59.mixer.dt_bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.59.mixer.in_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.59.mixer.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.59.mixer.out_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.59.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.6.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.6.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.6.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.6.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.6.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.6.mixer.in_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.6.mixer.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.6.mixer.out_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.6.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.60.mixer.A_log": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.60.mixer.D": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.60.mixer.conv1d.bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.60.mixer.conv1d.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.60.mixer.dt_bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.60.mixer.in_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.60.mixer.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.60.mixer.out_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.60.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.61.mixer.A_log": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.61.mixer.D": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.61.mixer.conv1d.bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.61.mixer.conv1d.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.61.mixer.dt_bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.61.mixer.in_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.61.mixer.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.61.mixer.out_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.61.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.62.mixer.A_log": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.62.mixer.D": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.62.mixer.conv1d.bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.62.mixer.conv1d.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.62.mixer.dt_bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.62.mixer.in_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.62.mixer.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.62.mixer.out_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.62.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.63.mixer.A_log": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.63.mixer.D": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.63.mixer.conv1d.bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.63.mixer.conv1d.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.63.mixer.dt_bias": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.63.mixer.ib_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.63.mixer.in_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.63.mixer.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.63.mixer.out_proj.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.63.norm.weight": "pytorch_model-00006-of-00006.bin",
+    "backbone.layers.7.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.7.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.7.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.7.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.7.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.7.mixer.in_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.7.mixer.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.7.mixer.out_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.7.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.8.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.8.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.8.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.8.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.8.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.8.mixer.in_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.8.mixer.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.8.mixer.out_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.8.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.9.mixer.A_log": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.9.mixer.D": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.9.mixer.conv1d.bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.9.mixer.conv1d.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.9.mixer.dt_bias": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.9.mixer.in_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.9.mixer.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.9.mixer.out_proj.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.layers.9.norm.weight": "pytorch_model-00001-of-00006.bin",
+    "backbone.norm_f.weight": "pytorch_model-00006-of-00006.bin",
+    "lm_head.weight": "pytorch_model-00006-of-00006.bin"
+  }
+}

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59f95e28944c062244741268596badc900df86c7f5ded05088d2da22a7379e06
+size 587583