modeling_flame.py

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,
        )