feat: 'main'

__init__.py

@@ -0,0 +1,7 @@
+'''
+* @author: EmpyreanMoon
+*
+* @create: 2025-07-17 19:20
+*
+* @description: FLAME SMALL (2M)
+'''

config.json

@@ -0,0 +1,29 @@
+{
+  "activation": "gelu",
+  "architectures": [
+    "FLAMEForPrediction"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_flame.FLAMEConfig",
+    "AutoModelForCausalLM": "modeling_flame.FLAMEForPrediction"
+  },
+  "couple_layers": 3,
+  "d_conv": 4,
+  "d_couple": 256,
+  "d_ff": 512,
+  "d_model": 256,
+  "dec_layers": 1,
+  "dropout": 0.2,
+  "enc_layers": 1,
+  "expand": 2,
+  "head_dim": 64,
+  "head_dropout": 0.2,
+  "learnable": false,
+  "model_type": "flame",
+  "n_heads": 4,
+  "norm_mode": "layer",
+  "patch_len": 48,
+  "stride": 48,
+  "torch_dtype": "float32",
+  "transformers_version": "4.53.2"
+}

configuration_flame.py

@@ -0,0 +1,46 @@
+from transformers import PretrainedConfig
+
+
+class FLAMEConfig(PretrainedConfig):
+    model_type = "flame"
+
+    def __init__(
+            self,
+            patch_len: int = 48,
+            expand: int = 2,
+            d_conv: int = 4,
+            d_model: int = 256,
+            d_ff: int = 512,
+            d_couple: int = 256,
+            enc_layers: int = 1,
+            dec_layers: int = 3,
+            couple_layers: int = 5,
+            head_dim: int = 64,
+            n_heads: int = 4,
+            activation: str = "gelu",
+            dropout: float = 0.2,
+            head_dropout: float = 0.2,
+            learnable: bool = False,
+            norm_mode: str = 'layer',
+            **kwargs,
+    ):
+        self.patch_len = patch_len
+        self.expand = expand
+        self.d_conv = d_conv
+        self.d_model = d_model
+        self.d_ff = d_ff
+        self.d_couple = d_couple
+        self.enc_layers = enc_layers
+        self.dec_layers = dec_layers
+        self.couple_layers = couple_layers
+        self.head_dim = head_dim
+        self.n_heads = n_heads
+        self.activation = activation
+        self.dropout = dropout
+        self.head_dropout = head_dropout
+        self.learnable = learnable
+        self.norm_mode = norm_mode
+
+        super().__init__(
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.53.2"
+}

model.safetensors

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

modeling_flame.py

@@ -0,0 +1,828 @@
+from typing import Optional
+from typing import Tuple
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import zuko
+from einops import rearrange
+from mamba_ssm import Mamba2
+from torch import Tensor
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+
+from .configuration_flame import FLAMEConfig
+from .ts_generation_mixin import TSGenerationMixin
+
+
+class Transpose(nn.Module):
+    def __init__(self, *dims, contiguous=False):
+        super().__init__()
+        self.dims, self.contiguous = dims, contiguous
+
+    def forward(self, x):
+        if self.contiguous:
+            return x.transpose(*self.dims).contiguous()
+        else:
+            return x.transpose(*self.dims)
+
+
+class MultiheadAttention(nn.Module):
+    def __init__(self, d_model, n_heads, d_k=None, d_v=None, res_attention=False, attn_dropout=0., proj_dropout=0.,
+                 qkv_bias=True, lsa=False, rope_type=False):
+        """Multi Head Attention Layer
+        Input shape:
+            Q:       [batch_size (bs) x max_q_len x d_model]
+            K, V:    [batch_size (bs) x q_len x d_model]
+            mask:    [q_len x q_len]
+        """
+        super().__init__()
+        d_k = d_model // n_heads if d_k is None else d_k
+        d_v = d_model // n_heads if d_v is None else d_v
+
+        self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v
+
+        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
+        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
+        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=qkv_bias)
+
+        # Scaled Dot-Product Attention (multiple heads)
+        self.res_attention = res_attention
+        self.sdp_attn = ScaledDotProductAttention(d_model, n_heads, attn_dropout=attn_dropout,
+                                                  res_attention=self.res_attention, lsa=lsa, rope_type=rope_type)
+
+        # Poject output
+        self.to_out = nn.Sequential(nn.Linear(n_heads * d_v, d_model), nn.Dropout(proj_dropout))
+
+    def forward(self, Q: Tensor, K: Optional[Tensor] = None, V: Optional[Tensor] = None, prev: Optional[Tensor] = None,
+                key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None):
+
+        bs = Q.size(0)
+        if K is None: K = Q
+        if V is None: V = Q
+
+        # Linear (+ split in multiple heads)
+        q_s = self.W_Q(Q).view(bs, -1, self.n_heads, self.d_k).transpose(1,
+                                                                         2)  # q_s    : [bs x n_heads x max_q_len x d_k]
+        k_s = self.W_K(K).view(bs, -1, self.n_heads, self.d_k).permute(0, 2, 3,
+                                                                       1)  # k_s    : [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)
+        v_s = self.W_V(V).view(bs, -1, self.n_heads, self.d_v).transpose(1, 2)  # v_s    : [bs x n_heads x q_len x d_v]
+
+        # Apply Scaled Dot-Product Attention (multiple heads)
+        if self.res_attention:
+            output, attn_weights, attn_scores = self.sdp_attn(q_s, k_s, v_s, prev=prev,
+                                                              key_padding_mask=key_padding_mask, attn_mask=attn_mask)
+        else:
+            output, attn_weights = self.sdp_attn(q_s, k_s, v_s, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
+        # output: [bs x n_heads x q_len x d_v], attn: [bs x n_heads x q_len x q_len], scores: [bs x n_heads x max_q_len x q_len]
+
+        # back to the original inputs dimensions
+        output = output.transpose(1, 2).contiguous().view(bs, -1,
+                                                          self.n_heads * self.d_v)  # output: [bs x q_len x n_heads * d_v]
+        output = self.to_out(output)
+
+        if self.res_attention:
+            return output, attn_weights, attn_scores
+        else:
+            return output, attn_weights
+
+
+class ScaledDotProductAttention(nn.Module):
+    r"""Scaled Dot-Product Attention module (Attention is all you need by Vaswani et al., 2017) with optional residual attention from previous layer
+    (Realformer: Transformer likes residual attention by He et al, 2020) and locality self sttention (Vision Transformer for Small-Size Datasets
+    by Lee et al, 2021)"""
+
+    def __init__(self, d_model, n_heads, attn_dropout=0., res_attention=False, lsa=False, rope_type=False):
+        super().__init__()
+        self.attn_dropout = nn.Dropout(attn_dropout)
+        self.res_attention = res_attention
+        head_dim = d_model // n_heads
+        self.scale = nn.Parameter(torch.tensor(head_dim ** -0.5), requires_grad=lsa)
+        self.lsa = lsa
+        self.rope_type = rope_type
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor, prev: Optional[Tensor] = None,
+                key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None):
+        '''
+        Input shape:
+            q               : [bs x n_heads x max_q_len x d_k]
+            k               : [bs x n_heads x d_k x seq_len]
+            v               : [bs x n_heads x seq_len x d_v]
+            prev            : [bs x n_heads x q_len x seq_len]
+            key_padding_mask: [bs x seq_len]
+            attn_mask       : [1 x seq_len x seq_len]
+        Output shape:
+            output:  [bs x n_heads x q_len x d_v]
+            attn   : [bs x n_heads x q_len x seq_len]
+            scores : [bs x n_heads x q_len x seq_len]
+        '''
+        # using RoPE
+        if self.rope_type:
+            q, k = RoPE_decoder(q, k.permute(0, 1, 3, 2))
+        else:
+            q, k = RoPE(q, k.permute(0, 1, 3, 2))
+        k = k.permute(0, 1, 3, 2)
+
+        # Scaled MatMul (q, k) - similarity scores for all pairs of positions in an input sequence
+        attn_scores = torch.matmul(q, k) * self.scale  # attn_scores : [bs x n_heads x max_q_len x q_len]
+
+        # Add pre-softmax attention scores from the previous layer (optional)
+        if prev is not None: attn_scores = attn_scores + prev
+
+        # Attention mask (optional)
+        if attn_mask is not None:  # attn_mask with shape [q_len x seq_len] - only used when q_len == seq_len
+            if attn_mask.dtype == torch.bool:
+                attn_scores.masked_fill_(attn_mask, -np.inf)
+            else:
+                attn_scores += attn_mask
+
+        # Key padding mask (optional)
+        if key_padding_mask is not None:  # mask with shape [bs x q_len] (only when max_w_len == q_len)
+            attn_scores.masked_fill_(key_padding_mask.unsqueeze(1).unsqueeze(2), -np.inf)
+
+        # normalize the attention weights
+        attn_weights = F.softmax(attn_scores, dim=-1)  # attn_weights   : [bs x n_heads x max_q_len x q_len]
+        attn_weights = self.attn_dropout(attn_weights)
+
+        # compute the new values given the attention weights
+        output = torch.matmul(attn_weights, v)  # output: [bs x n_heads x max_q_len x d_v]
+
+        if self.res_attention:
+            return output, attn_weights, attn_scores
+        else:
+            return output, attn_weights
+
+
+def RoPE(q, k):
+    # q,k: (bs, head, max_len, output_dim)
+    batch_size = q.shape[0]
+    nums_head = q.shape[1]
+    max_len = q.shape[2]
+    output_dim = q.shape[-1]
+
+    # (bs, head, max_len, output_dim)
+    pos_emb = sinusoidal_position_embedding(batch_size, nums_head, max_len, output_dim, q.device, factor=1)
+
+    # cos_pos,sin_pos: (bs, head, max_len, output_dim)
+    # 看rope公式可知，相邻cos，sin之间是相同的，所以复制一遍。如(1,2,3)变成(1,1,2,2,3,3)
+    cos_pos = pos_emb[..., 1::2].repeat_interleave(2, dim=-1)  # 将奇数列信息抽取出来也就是cos 拿出来并复制
+    sin_pos = pos_emb[..., ::2].repeat_interleave(2, dim=-1)  # 将偶数列信息抽取出来也就是sin 拿出来并复制
+
+    # q,k: (bs, head, max_len, output_dim)
+    q2 = torch.stack([-q[..., 1::2], q[..., ::2]], dim=-1)
+    q2 = q2.reshape(q.shape)  # reshape后就是正负交替了
+
+    # 更新qw, *对应位置相乘
+    q = q * cos_pos + q2 * sin_pos
+
+    k2 = torch.stack([-k[..., 1::2], k[..., ::2]], dim=-1)
+    k2 = k2.reshape(k.shape)
+    # 更新kw, *对应位置相乘
+    k = k * cos_pos + k2 * sin_pos
+
+    return q, k
+
+
+def RoPE_decoder(q, k):
+    # q,k: (bs, head, max_len, output_dim)
+    batch_size = q.shape[0]
+    nums_head = q.shape[1]
+    q_max_len = q.shape[2]
+    k_max_len = k.shape[2]
+    output_dim = q.shape[-1]
+
+    # (bs, head, max_len, output_dim)
+    pos_emb = sinusoidal_position_embedding(batch_size, nums_head, k_max_len + q_max_len, output_dim, q.device,
+                                            factor=1)
+
+    # cos_pos,sin_pos: (bs, head, max_len, output_dim)
+    # 看rope公式可知，相邻cos，sin之间是相同的，所以复制一遍。如(1,2,3)变成(1,1,2,2,3,3)
+    cos_pos = pos_emb[..., 1::2].repeat_interleave(2, dim=-1)  # 将奇数列信息抽取出来也就是cos 拿出来并复制
+    sin_pos = pos_emb[..., ::2].repeat_interleave(2, dim=-1)  # 将偶数列信息抽取出来也就是sin 拿出来并复制
+
+    # q,k: (bs, head, max_len, output_dim)
+    q2 = torch.stack([-q[..., 1::2], q[..., ::2]], dim=-1)
+    q2 = q2.reshape(q.shape)  # reshape后就是正负交替了
+
+    # 更新qw, *对应位置相乘
+    q = q * cos_pos[:, :, -q_max_len:, :] + q2 * sin_pos[:, :, -q_max_len:, :]
+
+    k2 = torch.stack([-k[..., 1::2], k[..., ::2]], dim=-1)
+    k2 = k2.reshape(k.shape)
+    # 更新kw, *对应位置相乘
+    k = k * cos_pos[:, :, :k_max_len, :] + k2 * sin_pos[:, :, :k_max_len, :]
+    return q, k
+
+
+def sinusoidal_position_embedding(batch_size, nums_head, max_len, output_dim, device, factor=1.0):
+    # (max_len * factor, 1)
+    position = torch.arange(0, max_len * factor, 1 / factor, dtype=torch.float).unsqueeze(-1)
+    # (output_dim//2)
+    ids = torch.arange(0, output_dim // 2, dtype=torch.float)  # i 范围是 [0, d/2]
+    theta = torch.pow(10000, -2 * ids / output_dim)
+
+    # (max_len * factor, output_dim//2)
+    embeddings = position * theta
+
+    # (max_len * factor, output_dim//2, 2)
+    embeddings = torch.stack([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)
+
+    # (bs, head, max_len * factor, output_dim//2, 2)
+    embeddings = embeddings.repeat((batch_size, nums_head, *([1] * len(embeddings.shape))))
+
+    # (bs, head, max_len * factor, output_dim)
+    embeddings = torch.reshape(embeddings, (batch_size, nums_head, -1, output_dim))
+    embeddings = embeddings.to(device)
+
+    # 如果 factor > 1, 使用插值位置来生成更细粒度的嵌入
+    if factor > 1.0:
+        interpolation_indices = torch.linspace(0, embeddings.shape[2] - 1, max_len).long()
+        embeddings = embeddings[:, :, interpolation_indices, :]
+
+    return embeddings
+
+
+def causal_attention_mask(seq_length):
+    mask = torch.triu(torch.ones(seq_length, seq_length) * float('-inf'), diagonal=1)
+    return mask.unsqueeze(0).unsqueeze(0)
+
+
+def resize(x_tensor, new_shape):
+    return F.interpolate(x_tensor.unsqueeze(0), size=new_shape, mode='linear').squeeze(0)
+
+
+def resample(old: torch.Tensor, new_patch_len: int):
+    assert old.dim() == 2, "the size of input tensor should be (d_model, patch_size)"
+    if old.size(1) == new_patch_len:
+        return old
+
+    old = old.T
+    old_shape = old.size(0)
+    factor = new_patch_len / old_shape
+
+    basis_vectors = torch.eye(old_shape, dtype=torch.get_default_dtype(), device=old.device)
+    resize_mat = resize(basis_vectors, new_patch_len).T
+    resize_mat_pinv = torch.linalg.pinv(resize_mat.T)
+
+    resampled_kernels = resize_mat_pinv @ old * math.sqrt(factor)
+
+    return resampled_kernels.T
+
+
+class MambaDecoder(nn.Module):
+    def __init__(self, configs):
+        super(MambaDecoder, self).__init__()
+        self.mamba_dec = nn.ModuleList(
+            [DecoderLayer(configs)
+             for _ in range(configs.dec_layers)])
+
+    def forward(self, x_enc, x_rec):
+        x_dec = x_enc
+        for layer in self.mamba_dec:
+            x_dec = layer(x_dec, x_rec)
+
+        return x_dec
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, configs):
+        super(DecoderLayer, self).__init__()
+        self.mamba = Mamba2(d_model=configs.d_model,
+                            expand=configs.expand,
+                            d_state=configs.d_ff,
+                            d_conv=configs.d_conv,
+                            headdim=configs.head_dim)
+
+        self.cross_attention = MultiheadAttention(configs.d_model, configs.n_heads, attn_dropout=configs.dropout,
+                                                  rope_type=True)
+        self.mlp = nn.Sequential(nn.Linear(configs.d_model, configs.d_ff), nn.SiLU(),
+                                 nn.Linear(configs.d_ff, configs.d_model))
+
+        if configs.norm_mode == 'batch':
+            self.norm1 = nn.Sequential(Transpose(1, 2), nn.BatchNorm1d(configs.d_model), Transpose(1, 2))
+            self.norm2 = nn.Sequential(Transpose(1, 2), nn.BatchNorm1d(configs.d_model), Transpose(1, 2))
+            self.norm3 = nn.Sequential(Transpose(1, 2), nn.BatchNorm1d(configs.d_model), Transpose(1, 2))
+        else:
+            self.norm1 = nn.LayerNorm(configs.d_model)
+            self.norm2 = nn.LayerNorm(configs.d_model)
+            self.norm3 = nn.LayerNorm(configs.d_model)
+
+    def forward(self, x_enc, x_rec):
+        x_dec = self.mamba(x_enc)
+        x_dec = self.norm1(x_dec) + x_enc
+
+        tokens, _ = self.cross_attention(x_dec, x_rec, x_rec)
+        tokens = self.norm2(tokens) + x_dec
+
+        repr = self.mlp(tokens)
+        repr = self.norm3(repr) + tokens
+
+        return repr
+
+
+class TSTEncoder(nn.Module):
+    def __init__(self, configs, norm='BatchNorm', activation='gelu', res_attention=False, pre_norm=False,
+                 store_attn=False):
+        super().__init__()
+
+        self.layers = nn.ModuleList(
+            [TSTEncoderLayer(configs.d_model, n_heads=configs.n_heads, d_ff=configs.d_ff, norm=norm,
+                             attn_dropout=configs.dropout, dropout=configs.head_dropout,
+                             activation=activation, res_attention=res_attention,
+                             pre_norm=pre_norm, store_attn=store_attn) for _ in
+             range(configs.enc_layers)])
+        self.res_attention = res_attention
+
+    def forward(self, src: Tensor):
+        """
+        src: tensor [bs x q_len x d_model]
+        """
+        output = src
+        scores = None
+        if self.res_attention:
+            for mod in self.layers: output, scores = mod(output, prev=scores)
+            return output
+        else:
+            for mod in self.layers: output = mod(output)
+            return output
+
+
+class TSTEncoderLayer(nn.Module):
+    def __init__(self, d_model, n_heads, d_ff=256, store_attn=False,
+                 norm='LayerNorm', attn_dropout=0, dropout=0., bias=True,
+                 activation="gelu", res_attention=False, pre_norm=False):
+        super().__init__()
+        assert not d_model % n_heads, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
+        d_k = d_model // n_heads
+        d_v = d_model // n_heads
+
+        # Multi-Head attention
+        self.res_attention = res_attention
+        self.self_attn = MultiheadAttention(d_model, n_heads, d_k, d_v, attn_dropout=attn_dropout, proj_dropout=dropout,
+                                            res_attention=res_attention)
+
+        # Add & Norm
+        self.dropout_attn = nn.Dropout(dropout)
+        if "batch" in norm.lower():
+            self.norm_attn = nn.Sequential(Transpose(1, 2), nn.BatchNorm1d(d_model), Transpose(1, 2))
+        else:
+            self.norm_attn = nn.LayerNorm(d_model)
+
+        # Position-wise Feed-Forward
+        self.ff = nn.Sequential(nn.Linear(d_model, d_ff, bias=bias),
+                                get_activation_fn(activation),
+                                nn.Dropout(dropout),
+                                nn.Linear(d_ff, d_model, bias=bias))
+
+        # Add & Norm
+        self.dropout_ffn = nn.Dropout(dropout)
+        if "batch" in norm.lower():
+            self.norm_ffn = nn.Sequential(Transpose(1, 2), nn.BatchNorm1d(d_model), Transpose(1, 2))
+        else:
+            self.norm_ffn = nn.LayerNorm(d_model)
+
+        self.pre_norm = pre_norm
+        self.store_attn = store_attn
+
+        # # se block
+        # self.SE = SE_Block(inchannel=7)
+
+    def forward(self, src: Tensor, prev: Optional[Tensor] = None):
+        """
+        src: tensor [bs x q_len x d_model]
+        """
+        # Multi-Head attention sublayer
+        if self.pre_norm:
+            src = self.norm_attn(src)
+        ## Multi-Head attention
+        if self.res_attention:
+            src2, attn, scores = self.self_attn(src, src, src, prev)
+        else:
+            # attention_mask = causal_attention_mask(src.shape[1]).to(src.device)
+            # src2, attn = self.self_attn(src, src, src, attn_mask=attention_mask)
+            src2, attn = self.self_attn(src, src, src)
+        if self.store_attn:
+            self.attn = attn
+
+        # total, num_patch, d_model = src2.size()
+        # bs = int(total/7)
+
+        # src2 = self.SE(src2.reshape(bs, 7, num_patch, -1)).reshape(total, num_patch, -1)
+
+        ## Add & Norm
+        src = src + self.dropout_attn(src2)  # Add: residual connection with residual dropout
+        if not self.pre_norm:
+            src = self.norm_attn(src)
+
+        # Feed-forward sublayer
+        if self.pre_norm:
+            src = self.norm_ffn(src)
+        ## Position-wise Feed-Forward
+        src2 = self.ff(src)
+        ## Add & Norm
+        src = src + self.dropout_ffn(src2)  # Add: residual connection with residual dropout
+        if not self.pre_norm:
+            src = self.norm_ffn(src)
+
+        if self.res_attention:
+            return src, scores
+        else:
+            return src
+
+
+def get_activation_fn(activation):
+    if callable(activation):
+        return activation()
+    elif activation.lower() == "relu":
+        return nn.ReLU()
+    elif activation.lower() == "gelu":
+        return nn.GELU()
+    raise ValueError(f'{activation} is not available. You can use "relu", "gelu", or a callable')
+
+
+class PatchEmbedding(nn.Module):
+    def __init__(self, configs):
+        super(PatchEmbedding, self).__init__()
+        self.patch_len = configs.patch_len
+        self.d_model = configs.d_model
+        self.proj = nn.Linear(self.patch_len, self.d_model, bias=False)
+
+    def forward(self, x):
+        output = self.proj(x)
+        return output
+
+
+class LegendreMemory(nn.Module):
+    def __init__(self, configs):
+        super(LegendreMemory, self).__init__()
+        self.d_model = configs.d_model
+        A, B = self._gen_AB_base_matrices(self.d_model)
+        if configs.learnable:
+            self.A = nn.Parameter(A)
+            self.B = nn.Parameter(B)
+        else:
+            self.register_buffer("A", A)
+            self.register_buffer("B", B)
+
+    def _pytorch_cont2discrete_zoh(
+            self, A: torch.Tensor, B: torch.Tensor, dt: float = 1.0
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+            Pytorch-specific implementation of discretization of a continuous-time
+            state space model using zero-order hold (ZOH) on the inputs.
+        """
+        em_upper = torch.cat((A, B), dim=1)
+        # Need to stack zeros under the a and b matrices
+        em_lower = torch.cat((
+            torch.zeros((B.shape[1], B.shape[0]), dtype=A.dtype, device=A.device),
+            torch.zeros((B.shape[1], B.shape[1]), dtype=A.dtype, device=A.device)
+        ), dim=1)
+
+        em = torch.cat((em_upper, em_lower), dim=0)
+        ms = torch.linalg.matrix_exp(dt * em)
+
+        # Dispose of the lower rows
+        ms = ms[:A.shape[0], :]
+
+        ad = ms[:, :A.shape[1]]
+        bd = ms[:, A.shape[1]:]
+
+        return ad, bd
+
+    def _gen_AB_base_matrices(self, order: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Compute analog A/B matrices
+        Q = torch.arange(order, dtype=torch.float64)
+        R = (2 * Q + 1).unsqueeze(1)
+        i, j = torch.meshgrid(Q, Q, indexing="ij")
+        A = torch.where(i < j, -1, (-1.0) ** (i - j + 1)) * R
+        B = (-1.0) ** Q.unsqueeze(1) * R
+        return A, B
+
+    def _gen_AB(self, theta, dt=1.0) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Discretize
+        Ad, Bd = self._pytorch_cont2discrete_zoh(self.A / theta, self.B / theta, dt)
+        return Ad.float(), Bd.float()
+
+    def _one_step(self, Ad, Bd, m, u):
+        m_prime = torch.einsum('dk,bk -> bd', Ad, m) + torch.einsum('dk,bk->bd', Bd, u)
+        return m_prime
+
+    def forward(self, x):
+        b, theta = x.shape
+        Ad, Bd = self._gen_AB(theta)
+
+        m = torch.zeros(b, self.d_model, dtype=x.dtype, device=x.device)
+        for i in range(theta):
+            u = x[:, i:i + 1]
+            m = self._one_step(Ad, Bd, m, u)
+
+        return m
+
+
+class TimeDelayEmbedding(nn.Module):
+    def __init__(self, configs):
+        super(TimeDelayEmbedding, self).__init__()
+        self.patch_len = configs.patch_len
+        self.stride = configs.stride
+        self.Legendre_Memory = LegendreMemory(configs)
+
+    def _period_search(self, x):
+        xf = torch.fft.rfft(x, dim=-1)
+        # find period by amplitudes
+        frequency_list = abs(xf).mean(0)
+        frequency_list[0] = 0
+        _, top_list = torch.topk(frequency_list, 1)
+        top_list = top_list.detach().cpu().numpy()
+        period = x.shape[1] // top_list
+        return period
+
+    def _embedding(self, x, period=None):
+        if period is None:
+            period = list(self._period_search(x))[0]
+        patch_len = period
+        patches = x.unfold(dimension=-1, size=patch_len, step=patch_len)
+        b, n, p = patches.shape
+        patches = rearrange(patches, 'b n p -> (b n) p')
+
+        embedding = self.Legendre_Memory(patches)
+        embedding = rearrange(embedding, '(b n) d -> b n d', b=b, n=n)
+        return embedding
+
+    def forward(self, x, period=None, num_samples=1):
+        if not self.training and num_samples == 1:
+            return self._embedding(x, period=period)
+
+        if period is None:
+            period = list(self._period_search(x))[0]
+
+        seq_len = x.shape[-1]
+        period = period if period < seq_len else self.patch_len
+
+        padding_num = ((self.patch_len + period - 1) // period) * period - self.patch_len
+        padding_action = nn.ReplicationPad1d((padding_num, 0))
+        padded_x = padding_action(x)
+        new_patch_len = self.patch_len + padding_num
+        patches = padded_x.unfold(dimension=-1, size=new_patch_len, step=self.stride)
+        b, n, p = patches.shape
+        patches = rearrange(patches, 'b n p -> (b n) p')
+
+        embedding = self.Legendre_Memory(patches)
+        embedding = rearrange(embedding, '(b n) d -> b n d', b=b, n=n)
+        return embedding
+
+
+class MixedEmbedding(nn.Module):
+    def __init__(self, configs):
+        super(MixedEmbedding, self).__init__()
+        self.configs = configs
+        self.patch_len = configs.patch_len
+        self.stride = configs.patch_len
+        self.d_model = configs.d_model
+
+        # TimeDelayEmbedding
+        self.time_delay_embedding = TimeDelayEmbedding(configs)
+
+        # PatchEmbedding
+        self.patch_embedding = PatchEmbedding(configs)
+
+        self.dropout = nn.Dropout(configs.dropout)
+
+    def _flex_embedding(self, x, inference_patch_len):
+        patch_len = inference_patch_len
+        seq_len = x.shape[-1]
+        patch_num = math.ceil((seq_len - patch_len) / patch_len) + 1
+        padding = patch_num * patch_len - seq_len
+        padding_patch_layer = nn.ReplicationPad1d((0, padding))
+        x = padding_patch_layer(x)
+
+        # [batch_size, patch_num, patch_size]
+        patches = x.unfold(dimension=-1, size=patch_len, step=patch_len)
+
+        proj = nn.Linear(patch_len, self.d_model, bias=False)
+        proj.weight.data = resample(old=self.patch_embedding.proj.weight.data, new_patch_len=patch_len)
+
+        patch_embedding = proj(patches)
+        time_delay_embedding = self.time_delay_embedding(x, period=patch_len, num_samples=1)
+        embedding = patch_embedding + time_delay_embedding
+        return embedding
+
+    def forward(self, x, inference_patch_len=48, num_samples=1):
+        # do patching
+        # padding for the original stride
+        if not self.training and num_samples == 1:
+            return self._flex_embedding(x, inference_patch_len)
+
+        seq_len = x.shape[-1]
+        patch_num = math.ceil((seq_len - self.patch_len) / self.stride) + 1
+        padding = self.patch_len + (patch_num - 1) * self.stride - seq_len
+        padding_patch_layer = nn.ReplicationPad1d((0, padding))
+        x = padding_patch_layer(x)
+
+        # [batch_size, patch_num, patch_size]
+        patches = x.unfold(dimension=-1, size=self.patch_len, step=self.stride)
+
+        # [batch_size, patch_num, d_model]
+        patch_embedding = self.patch_embedding(patches)
+
+        # [batch_size, patch_num, d_model]
+        time_delay_embedding = self.time_delay_embedding(x, num_samples=num_samples)
+
+        embedding = patch_embedding + time_delay_embedding
+
+        return self.dropout(embedding)
+
+
+class FLAMEModel(nn.Module):
+    def __init__(self, configs):
+        super(FLAMEModel, self).__init__()
+        self.patch_len = configs.patch_len
+        configs.stride = configs.patch_len
+
+        self.embedding = MixedEmbedding(configs)
+
+        self.d_model = configs.d_model
+
+        self.encoder = TSTEncoder(configs)
+
+        self.decoder = MambaDecoder(configs)
+
+        self.proj = nn.Linear(configs.d_model, configs.patch_len, bias=False)
+        self.dropout = nn.Dropout(configs.head_dropout)
+
+        self.flow = zuko.flows.MAF(features=configs.patch_len, context=configs.d_model,
+                                   transforms=configs.couple_layers,
+                                   hidden_features=[configs.d_couple] * configs.couple_layers)
+        self.configs = configs
+
+    def _prob_head(self, dec_out):
+        tokens = rearrange(dec_out, 'b n d -> (b n) d')
+        dist = self.flow(tokens)
+        return dist
+
+    def _get_weights(self, n_preds, decay_rate=0.5):
+        """
+        Generate dynamic weights for the replicated tokens using an exponential decay scheme.
+
+        Args:
+        - n_preds (int): Number of predictions to generate weights for.
+        - decay_rate (float): The base of the exponential decay. Lower values decay faster (default: 0.9).
+
+        Returns:
+        - torch.Tensor: A tensor of weights with exponential decay.
+        """
+        # Exponential decay weights
+        weights = decay_rate ** torch.arange(n_preds)
+        return weights
+
+    def forward(self, input, target=None, pred_len=None, inference_patch_len=48, num_samples=1):
+        if not self.training:
+            return self._predict(input, pred_len=pred_len, inference_patch_len=inference_patch_len,
+                                 num_samples=num_samples)
+        else:
+            return self._loss(input=input, target=target)
+
+    def _loss(self, input, target, eps=1e2):
+        # forward
+        pred_len = input.shape[-1]
+        x_enc = self.embedding(input)
+
+        x_enc = self.encoder(x_enc)
+        x_rec = rearrange(x_enc, 'b n p -> b (n p)')
+
+        predict_token_num = math.ceil(pred_len / self.patch_len)
+        weights = self._get_weights(predict_token_num).unsqueeze(0).unsqueeze(-1).to(input.device)
+        last_token = x_enc[:, -1:, :]
+        x_enc = weights * last_token.repeat(1, predict_token_num, 1)
+        # decoding
+        x_dec = self.decoder(x_enc, x_rec)
+
+        dec_out = self.proj(self.dropout(x_dec))
+
+        point_forecasts = rearrange(dec_out, 'b n p -> b (n p)')
+
+        forecasts = point_forecasts[:, :pred_len]
+
+        dist = self._prob_head(x_dec.detach())
+
+        # calculate loss
+        point_loss = self.point_loss(forecasts, target)
+        rec_loss = self.point_loss(x_rec, input)
+        transformed_target = rearrange(target, 'b (n p) c -> (b c n) p', p=self.patch_len)
+        raw_prob_loss = -dist.log_prob(transformed_target)
+        mask = raw_prob_loss < eps
+        raw_prob_loss = torch.where(mask, raw_prob_loss, torch.zeros_like(raw_prob_loss))
+
+        prob_loss = raw_prob_loss.mean() / self.patch_len
+
+        return point_loss + rec_loss + prob_loss
+
+    def _predict(self, input, pred_len, inference_patch_len=48, num_samples=None):
+        if num_samples is not None and num_samples > 1:
+            return self._prob_predict(input, pred_len, num_samples)
+
+        x_enc = self.embedding(input, inference_patch_len, num_samples=1)
+
+        x_rec = self.encoder(x_enc)
+
+        predict_token_num = math.ceil(pred_len / inference_patch_len)
+        weights = self._get_weights(predict_token_num).unsqueeze(0).unsqueeze(-1).to(input.device)
+        last_token = x_rec[:, -1:, :]
+        x_enc = weights * last_token.repeat(1, predict_token_num, 1)
+        # decoding
+        x_dec = self.decoder(x_enc, x_rec)
+
+        proj = nn.Linear(self.d_model, inference_patch_len, bias=False)
+        proj.weight.data = resample(old=self.proj.weight.data.T,
+                                    new_patch_len=inference_patch_len).T
+
+        dec_out = proj(x_dec)
+        point_forecasts = rearrange(dec_out, 'b n p -> b (n p)')
+        return point_forecasts[:, :pred_len]
+
+    def _prob_predict(self, input, pred_len, num_samples=None):
+
+        x_enc = self.embedding(input, num_samples=num_samples)
+
+        x_rec = self.encoder(x_enc)
+
+        predict_token_num = math.ceil(pred_len / self.patch_len)
+        weights = self._get_weights(predict_token_num).unsqueeze(0).unsqueeze(-1).to(input.device)
+        last_token = x_rec[:, -1:, :]
+        x_enc = weights * last_token.repeat(1, predict_token_num, 1)
+        # decoding
+        x_dec = self.decoder(x_enc, x_rec)
+
+        dist = self._prob_head(x_dec)
+
+        samples = dist.sample((num_samples,))
+        samples = rearrange(samples, 's (b n) p -> b s (n p) ', n=predict_token_num)[:, :, :pred_len]
+
+        prob_forecasts = samples
+
+        return prob_forecasts
+
+
+class FLAMEPretrainedModel(PreTrainedModel):
+    config_class = FLAMEConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["TSTEncoder", "MambaDecoder"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = False
+    _supports_cache_class = False
+
+
+class FLAMEForPrediction(FLAMEPretrainedModel, TSGenerationMixin):
+    def __init__(self, config: FLAMEConfig):
+        super().__init__(config)
+        self.config = config
+        self.model = FLAMEModel(config)
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    def forward(
+            self,
+            input_ids: torch.FloatTensor = None,
+            labels: Optional[torch.FloatTensor] = None,
+            max_output_length: Optional[int] = None,
+            revin: Optional[bool] = True,
+            num_samples: Optional[int] = 1,
+            inference_patch_len: Optional[int] = 48,
+    ):
+        if revin:
+            means = input_ids.mean(1, keepdim=True).detach()
+            stdev = input_ids.std(dim=1, keepdim=True, unbiased=False).detach() + 1e-5
+            input_ids = (input_ids - means) / stdev
+
+        outputs = self.model(
+            input=input_ids,
+            target=labels,
+            inference_patch_len=inference_patch_len,
+            num_samples=num_samples,
+            pred_len=max_output_length
+        )
+
+        loss = None
+        if labels is not None:
+            loss = outputs
+        else:
+            forecasts = outputs
+
+            if forecasts.ndim == 2:
+                forecasts = forecasts.unsqueeze(1)
+                forecasts = forecasts.repeat(1, num_samples, 1)
+            if revin:
+                stdev = stdev.unsqueeze(1).repeat(1, num_samples, 1)
+                means = means.unsqueeze(1).repeat(1, num_samples, 1)
+                forecasts = (forecasts * stdev) + means
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            logits=forecasts,
+        )

ts_generation_mixin.py

@@ -0,0 +1,63 @@
+from typing import Any, Dict, List, Optional, Union, Callable
+
+import torch
+from transformers import GenerationMixin, LogitsProcessorList, StoppingCriteriaList
+from transformers.generation.utils import GenerationConfig, GenerateOutput
+from transformers.utils import ModelOutput
+
+
+class TSGenerationMixin(GenerationMixin):
+    @torch.no_grad()
+    def generate(
+            self,
+            inputs: Optional[torch.Tensor] = None,
+            generation_config: Optional[GenerationConfig] = None,
+            logits_processor: Optional[LogitsProcessorList] = None,
+            stopping_criteria: Optional[StoppingCriteriaList] = None,
+            prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
+            synced_gpus: Optional[bool] = None,
+            assistant_model: Optional["PreTrainedModel"] = None,
+            streamer: Optional["BaseStreamer"] = None,
+            negative_prompt_ids: Optional[torch.Tensor] = None,
+            negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+            revin: Optional[bool] = True,
+            num_samples: Optional[int] = 1,
+            max_output_length: Optional[int] = 96,
+            inference_patch_len: Optional[int] = 48,
+            **kwargs,
+    ) -> Union[GenerateOutput, torch.Tensor]:
+        if len(inputs.shape) != 2:
+            raise ValueError('Input shape must be: [batch_size, seq_len]')
+        if revin:
+            means = inputs.mean(dim=-1, keepdim=True)
+            stdev = inputs.std(dim=-1, keepdim=True, unbiased=False) + 1e-5
+            inputs = (inputs - means) / stdev
+
+        model_inputs = {
+            "input_ids": inputs,
+            "max_output_length": max_output_length,
+            "revin": False,
+            "num_samples": num_samples,
+            "inference_patch_len": inference_patch_len,
+        }
+
+        outputs = self(**model_inputs)
+
+        predictions = outputs.logits
+
+        if revin:
+            stdev = stdev.unsqueeze(1).repeat(1, num_samples, 1)
+            means = means.unsqueeze(1).repeat(1, num_samples, 1)
+            predictions = (predictions * stdev) + means
+
+        return predictions
+
+    def _update_model_kwargs_for_generation(
+            self,
+            outputs: ModelOutput,
+            model_kwargs: Dict[str, Any],
+            horizon_length: int = 1,
+            is_encoder_decoder: bool = False,
+            standardize_cache_format: bool = False,
+    ) -> Dict[str, Any]:
+        return model_kwargs