Update modeling_super_linear.py

modeling_super_linear.py

@@ -191,49 +191,79 @@ class NLinear(nn.Module):
         return x+seq_last # to [Batch, Output length, Channel]   
     
  
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
 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)
+    """
+    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, bias=False)
 
-    def forward(self, x):
-        # x: [Batch, Input length,Channel]
-        x_shape = x.shape
-        if len(x_shape) == 2:
-            x = x.unsqueeze(-1)
-        
-        B,L,V    = x.shape
-        if L <  self.seq_len:
-            in_features     = L
-            W               = self.Linear.weight.detach()
-            fixed_weights   = W[:, :L]
-            dynamic_weights = W[:, L:]
-
-            if in_features != W.size(1):
-                dynamic_weights = F.interpolate(dynamic_weights.unsqueeze(0).unsqueeze(0), size=(self.horizon, in_features-self.seq_len), mode='bilinear', align_corners=False).squeeze(0).squeeze(0)
-            if self.fixed_in != 0:
-                fixed_weights = F.interpolate(fixed_weights.unsqueeze(0).unsqueeze(0), size=(self.horizon, L), mode='bilinear', align_corners=False).squeeze(0).squeeze(0)
-
-            x = self.revin_layer(x, 'norm')
-            x = F.linear(x, torch.cat((fixed_weights, dynamic_weights), dim=1))
-            x = self.revin_layer(x, 'denorm')
-            if len(x_shape) == 2:
-                x = x.squeeze(-1)
-            return x
+        # your RevIN layer (must be defined elsewhere)
+        self.revin = RevIN(num_features=None, affine=False,
+                           norm_type=None, subtract_last=False)
 
-        
-        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')
+    # ------------------------------------------------------------------ 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)
+
+            # 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)
+
+        # ---------- RevIN denormalisation -------------------------------------
+        x = self.revin(x, "denorm")
+
+        if squeeze_last:
+            x = x.squeeze(-1)                   # (B,horizon)
+
+        return x
 
-        if len(x_shape) == 2:
-            x = x.squeeze(-1)
-        return x # to [Batch, Output length, Channel]  
 "-------------------------------------------------------------------------------------------------------------------"
 class SparseNoisyMoE(nn.Module):
     def __init__(self, configs, experts=None):