Update modeling_super_linear.py

modeling_super_linear.py

@@ -193,72 +193,67 @@ class NLinear(nn.Module):
  
 
 class RLinear(nn.Module):
-    """
-    Linear projection from a variable-length input (L) to a fixed horizon,
-    applied *independently per channel* and wrapped with RevIN.
-    """
-    def __init__(self, input_len: int, output_len: int):
-        super().__init__()
-        self.seq_len  = input_len          # “design” length
-        self.horizon  = output_len
-
-        # bias=False because you asked to drop the bias
-        self.linear   = nn.Linear(input_len, output_len)
-
-        # your RevIN layer (must be defined elsewhere)
-        self.revin = RevIN(num_features=None, affine=False,
-                           norm_type=None, subtract_last=False)
-
-    # ------------------------------------------------------------------ helpers
-    def _resize_weight(self, weight: torch.Tensor, new_in: int) -> torch.Tensor:
-        """
-        Bilinearly interpolate columns so the weight becomes (horizon, new_in).
-        """
-        if new_in == weight.shape[1]:
-            return weight                       # nothing to do
-        w4d = weight.unsqueeze(0).unsqueeze(0)  # (1,1,out,in)
-        w_resized = F.interpolate(
-            w4d,
-            size=(self.horizon, new_in),        # always ≥ 0, so no crash
-            mode="bilinear",
-            align_corners=False
-        )[0, 0]                                 # back to (out,new_in)
-        return w_resized
-
-    # ------------------------------------------------------------------ forward
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        x: (B,L,C) or (B,L)  →  (B,horizon,C) or (B,horizon)
-        """
-        squeeze_last = False
-        if x.dim() == 2:                        # (B,L)
-            x = x.unsqueeze(-1)                 # (B,L,1)
-            squeeze_last = True
-
-        B, L, C = x.shape
-
-        # ---------- RevIN normalisation ---------------------------------------
-        x = self.revin(x, "norm")
-
-        if L == self.seq_len:                   # fast path – no resizing
-            x = self.linear(x.permute(0, 2, 1))        # (B,C,horizon)
-            x = x.permute(0, 2, 1)                     # (B,horizon,C)
-
-        else:                                   # resize the weight once
-            W = self._resize_weight(self.linear.weight.detach(), L)  # (out,L)
+    def __init__(self, input_len, output_len):
+        super(RLinear, self).__init__()
+        self.Linear              = nn.Linear(input_len, output_len)
+        self.seq_len             = input_len
+        self.horizon             = output_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,mode):
+        if mode == 1:
+            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
+        else:
+            W = self.Linear.weight.detach()
+            W4d = W.unsqueeze(0).unsqueeze(0)            # (1, 1, out, in)
+
+            # resize   H → self.horizon   and   W → new_lookback
+            new_W = F.interpolate(
+                        W4d,
+                        size=(self.horizon, new_lookback),   # (H_out, W_out)
+                        mode='bilinear',
+                        align_corners=False
+                    )[0, 0]                                 # drop the two singleton dims
+
+            self.zero_shot_Linear = new_W                  # shape (self.horizon, new_lookback)
+        
 
-            # project each channel separately: (B,C,L) @ (L,out)ᵀ → (B,C,out)
-            x = x.permute(0, 2, 1)              # (B,C,L)
-            x = torch.matmul(x, W.t())          # (B,C,out)
-            x = x.permute(0, 2, 1)              # (B,horizon,C)
+    def forward(self, x):
+        # x: [Batch, Input length,Channel]
+        x_shape = x.shape
+        if x.shape[1] < self.seq_len:
+            #if self.zero_shot_Linear is None:
+                #print(F"new Lookkback : {x.shape[1]}")
+            
+            self.transform_model(x.shape[1],3)
+            x = x.clone()
+            #x = x * (x.shape[1]/512)
+            x = self.revin_layer(x, 'norm')
+            x = F.linear(x, self.zero_shot_Linear)
+            x = self.revin_layer(x, 'denorm')
+            #x = x * (512/x.shape[1])
+            return x
 
-        # ---------- RevIN denormalisation -------------------------------------
-        x = self.revin(x, "denorm")
+                
+        if len(x_shape) == 2:
+            x = x.unsqueeze(-1)
 
-        if squeeze_last:
-            x = x.squeeze(-1)                   # (B,horizon)
+        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')
 
-        return x
+        if len(x_shape) == 2:
+            x = x.squeeze(-1)
+        return x # to [Batch, Output length, Channel]  
 
 "-------------------------------------------------------------------------------------------------------------------"
 class SparseNoisyMoE(nn.Module):