Update modeling_super_linear.py

modeling_super_linear.py

@@ -194,19 +194,45 @@ class NLinear(nn.Module):
 class RLinear(nn.Module):
     def __init__(self, input_len, output_len):
         super(RLinear, self).__init__()
-        self.Linear = nn.Linear(input_len, output_len)
-        self.revin_layer = RevIN(num_features = None, affine=False, norm_type = None, subtract_last = False)
+        self.Linear              = nn.Linear(input_len, output_len)
+        self.seq_len             = input_len
+        self.revin_layer         = RevIN(num_features = None, affine=False, norm_type = None, subtract_last = False)
+        self.zero_shot_Linear    = None
+
+    def transform_model(self,new_lookback):
+        W              = self.Linear.weight.detach()
+        new_W          = W[:, -new_lookback:] 
+        original_norm  = torch.norm(W, p=2)
+        new_norm       = torch.norm(new_W, p=2)
+        final_scaling  = original_norm / new_norm if new_norm.item() != 0 else 1.0
+        new_W          = new_W * final_scaling
+
+        self.zero_shot_Linear        = new_W
+
 
     def forward(self, x):
         # x: [Batch, Input length,Channel]
         x_shape = x.shape
         if len(x_shape) == 2:
             x = x.unsqueeze(-1)
-        x = x.clone()
-        x = self.revin_layer(x, 'norm')
-        
-        x = self.Linear(x.permute(0,2,1)).permute(0,2,1).clone()
-        x = self.revin_layer(x, 'denorm')
+
+        B,L,V    = x.shape
+        if L <  self.seq_len and self.zero_shot_Linear is None:
+            print(F"New Lookback :{L}")
+            self.transform_model(L)
+
+        if L <  self.seq_len:
+            x = x.clone()
+            x = self.revin_layer(x, 'norm')
+            x = F.linear(x, self.zero_shot_Linear).unsqueeze_(-1)
+            x = self.revin_layer(x, 'denorm')
+
+        else:
+            x = x.clone()
+            x = self.revin_layer(x, 'norm')
+            x = self.Linear(x.permute(0,2,1)).permute(0,2,1).clone()
+            x = self.revin_layer(x, 'denorm')
+
         if len(x_shape) == 2:
             x = x.squeeze(-1)
         return x # to [Batch, Output length, Channel]