Update modeling_super_linear.py

modeling_super_linear.py

@@ -511,6 +511,10 @@ class superLinear(nn.Module):
 
 
     def forward(self, x_enc, x_mark_enc=None, x_dec=None, x_mark_dec=None, mask=None, freq=[None], get_prob=False):
+
+        if inf_pred_len is  None:
+            inf_pred_len = self.inf_pred_len
+
         if len(x_enc.shape) > 2:
             x = x_enc.permute(0, 2, 1)
             B, V, L = x.shape
@@ -518,8 +522,26 @@ class superLinear(nn.Module):
             x    = x_enc
             B, L = x.shape
             V    = 1
+
+        short_lookback = False
+        if L < self.seq_len:
+          #  print("test!")
+            #ceil - very bad heuristic!
+            scale_factor = self.seq_len / L
+            scale_factor = int(np.ceil(scale_factor))
+            orig_pred_len = inf_pred_len
+
+            inf_pred_len = inf_pred_len*scale_factor
+            x = interpolate(x_enc.permute(0, 2, 1), scale_factor=scale_factor, mode='linear')
+
+            x = x[:,: , -self.seq_len:]
+            orig_L = L
+            L = self.seq_len
+
+            short_lookback  = True
             
         x = x.reshape(B * V, L)
+        
         expert_probs = None
         
         if get_prob:
@@ -527,23 +549,22 @@ class superLinear(nn.Module):
         else:
             out, self.moe_loss = self.moe(x)
 
-
-        if self.auto_regressive and self.max_horizon < self.inf_pred_len:
-            #print("bitch")
+        if self.auto_regressive and self.max_horizon < inf_pred_len:
             outputs = [out]
             ar_x = torch.cat([x, out], dim=1)[:, -self.seq_len:]
-            for i in range(0, self.inf_pred_len, self.max_horizon):
+            for i in range(0, inf_pred_len, self.max_horizon):
                 ar_out, _ = self.moe(ar_x)
                 outputs.append(ar_out)
                 ar_x = torch.cat([ar_x, ar_out], dim=1)[:, -self.seq_len:]
-            out = torch.cat(outputs, dim=1)[:, :self.inf_pred_len]
-            
-        if len(x_enc.shape) > 2:
-            out = out.reshape(B, V, out.shape[-1])
-            result = out.permute(0, 2, 1)
-        else:
-            result =  out
+            out = torch.cat(outputs, dim=1)[:,:inf_pred_len]
+        out = out.reshape(B, V, out.shape[-1])
+        
 
+        if short_lookback:
+            out = interpolate(out, scale_factor=1/scale_factor, mode='linear')
+           # print(out.shape)
+            out  = out[:, :,:orig_pred_len]
+        result = out.permute(0, 2, 1)
         if get_prob:
             expert_probs = expert_probs.reshape(B, V, expert_probs.shape[-1])
             return result, expert_probs