Update modeling_super_linear.py

modeling_super_linear.py

@@ -263,6 +263,7 @@ class SparseNoisyMoE(nn.Module):
     def __init__(self, configs, experts=None):
         super(SparseNoisyMoE, self).__init__()
         input_dim = configs.seq_len
+        self.lookback = configs.seq_len
         output_dim = configs.pred_len
         self.k = configs.top_k_experts
         self.noise_std = configs.noisy_gating_std
@@ -327,6 +328,16 @@ class SparseNoisyMoE(nn.Module):
             
         return I
 
+
+    def fourier_interp_dim1(self,x, target_len: int = 512):
+
+        L = x.size(1)
+        X      = torch.fft.rfft(x, dim=1)                   # (..., 25, ...)
+        pad    = target_len // 2 + 1 - X.size(1)
+        X_pad  = torch.cat([X, X.new_zeros(*X.shape[:-1], pad)], dim=1)
+        y      = torch.fft.irfft(X_pad, n=target_len, dim=1)
+        return y
+
     def forward(self, x, get_prob=False):
         if self.use_fft:
             x_0 = self.get_periodogram(x, ker_len=self.ker_len, n=self.fft_len, con=self.con)
@@ -348,6 +359,8 @@ class SparseNoisyMoE(nn.Module):
         self.topk_gates = F.softmax(self.topk_values, dim=1)
         
         batch_size = x.size(0)
+        if x_enc.shape[1] < 512:
+            x = self.fourier_interp_dim1(x)
         expert_outputs = torch.stack([self.experts[i](x) for i in range(self.num_experts)], dim=1)
 
         topk_indices_expanded = topk_indices.unsqueeze(-1).expand(-1, -1, expert_outputs.size(2))
@@ -697,7 +710,8 @@ class SuperLinearForCausalLM(PreTrainedModel, GenerationMixin):
         
         if x_enc.shape[1] < 512:
             #x_enc = self.revin_layer(x_enc, 'norm')
-            x_enc = self.fourier_interp_dim1(x_enc)
+            #x_enc = self.fourier_interp_dim1(x_enc)
+            pass
 
         #self.backbone.inf_pred_len = 336