Update modeling_super_linear.py

modeling_super_linear.py

@@ -192,62 +192,48 @@ class NLinear(nn.Module):
     
  
 class RLinear(nn.Module):
-    """
-    Resizable linear projection from variable input length L to fixed horizon,
-    applied independently to every channel, **without bias**.
-    """
-    def __init__(self, input_len: int, output_len: int):
-        super().__init__()
-        self.seq_len  = input_len
-        self.horizon  = output_len
-
-        # ★ bias removed → bias=False
-        self.linear   = nn.Linear(input_len, output_len)
-
-        self.revin = RevIN(num_features=None, affine=False,
-                           norm_type=None, subtract_last=False)
-
-    @staticmethod
-    def _resize_weight(weight: torch.Tensor, new_in: int, horizon: int) -> torch.Tensor:
-        """Interpolate columns so weight becomes (horizon, new_in)."""
-        if new_in == weight.shape[1]:
-            return weight
-        w4d = weight.unsqueeze(0).unsqueeze(0)                 # (1,1,out,in)
-        w_resized = F.interpolate(
-            w4d, size=(horizon, new_in), mode="bilinear",
-            align_corners=False
-        )[0, 0]                                                # (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)
-        """
-        squeeze_last = False
-        if x.dim() == 2:                                       # (B,L)
-            x = x.unsqueeze(-1)                                # (B,L,1)
-            squeeze_last = True
-
-        B, L, C = x.shape
-        x = self.revin(x, "norm")
-
-        if L == self.seq_len:                                  # fast path
-            x = self.linear(x.permute(0, 2, 1))                # (B,C,horizon)
-            x = x.permute(0, 2, 1)                             # (B,horizon,C)
-        else:                                                  # resize path
-            W = self.linear.weight.detach()                    # (out,in)
-            W_resized = self._resize_weight(W, L, self.horizon)
-
-            # ★ bias removed → no "+ b"
-            x = x.permute(0, 2, 1)                             # (B,C,L)
-            x = torch.einsum("bcl,ol->bco", x, W_resized)      # (B,C,out)
-            x = x.permute(0, 2, 1)                             # (B,horizon,C)
-
-        x = self.revin(x, "denorm")
-        if squeeze_last:
-            x = x.squeeze(-1)
-        return x
+    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)
 
+
+    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   = self.weights[:, :L]
+            dynamic_weights = self.weights[:, L:]
+
+            if in_features != self.weights.size(1) or out_features != self.weights.size(0):
+                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, self.fixed_in), 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
+
+        
+        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]  
 "-------------------------------------------------------------------------------------------------------------------"
 class SparseNoisyMoE(nn.Module):
     def __init__(self, configs, experts=None):