Update modeling_super_linear.py

modeling_super_linear.py

@@ -192,71 +192,81 @@ 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.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)
-        
+    """
+    Resizable linear projection from variable input length L to fixed horizon.
+    Each channel is projected independently (no mixing across channels).
+    """
+    def __init__(self, input_len: int, output_len: int):
+        super().__init__()
+        self.seq_len = input_len          # “design” length
+        self.horizon = output_len
+        # plain weight + bias – we will *interpolate* weight when L ≠ seq_len
+        self.linear = nn.Linear(input_len, output_len, bias=True)
+
+        # RevIN — your own implementation; keep it stateless for safety
+        self.revin = RevIN(num_features=None, affine=False,
+                           norm_type=None, subtract_last=False)
+
+    def _resize_weight(self, weight: torch.Tensor, new_in: int) -> torch.Tensor:
+        """
+        Bilinearly interpolate the *columns* of `weight` so that
+        weight.shape becomes (horizon, new_in).
+        """
+        if new_in == weight.shape[1]:
+            return weight          # no resizing needed
+
+        # weight: (out, in)  →  (1,1,out,in)  so we can use interpolate
+        w4d = weight.unsqueeze(0).unsqueeze(0)
+        w_resized = F.interpolate(
+            w4d, size=(self.horizon, new_in), mode="bilinear",
+            align_corners=False
+        )[0, 0]                    # back to (out, new_in)
+        return w_resized
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: (B, L, C)  or  (B, L)  ⇒  (B, horizon, C)  or  (B, horizon)
+        """
+        # make sure x is 3-D
+        squeeze_last_dim = False
+        if x.dim() == 2:           # (B,L)
+            x = x.unsqueeze(-1)    # (B,L,1)
+            squeeze_last_dim = True
 
+        B, L, C = x.shape
 
-            
+        # RevIN normalisation (over time dimension)
+        x = self.revin(x, "norm")
 
-    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)
+        if L == self.seq_len:
+            # fast path — built-in linear works
+            # reshape so that each channel is treated independently
+            x = x.permute(0, 2, 1)             # (B,C,L)
+            x = self.linear(x)                 # (B,C,horizon)
+            x = x.permute(0, 2, 1)             # (B,horizon,C)
 
-            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
+        else:
+            # slow path — resize the weight to match L
+            # freeze current weight & bias
+            W = self.linear.weight.detach()    # (out,in)
+            b = self.linear.bias.detach()      # (out)
 
-                
-        if len(x_shape) == 2:
-            x = x.unsqueeze(-1)
+            W_resized = self._resize_weight(W, L)          # (out,L)
 
-        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')
+            # apply per channel
+            #   x: (B,L,C)  →  (B,C,L)   so that last dim is L
+            x = x.permute(0, 2, 1)
+            #   out = x @ W_resized.T + b
+            out = torch.einsum("bcl,ol->bco", x, W_resized) + b
+            x = out.permute(0, 2, 1)           # back to (B,horizon,C)
 
-        if len(x_shape) == 2:
-            x = x.squeeze(-1)
-        return x # to [Batch, Output length, Channel]  
+        # RevIN denorm
+        x = self.revin(x, "denorm")
+
+        if squeeze_last_dim:
+            x = x.squeeze(-1)      # (B,horizon)
+
+        return x
 
 
 "-------------------------------------------------------------------------------------------------------------------"