Update modeling_super_linear.py

modeling_super_linear.py

@@ -293,41 +293,51 @@ class SparseNoisyMoE(nn.Module):
             self.gating_network = nn.Linear(input_dim, self.num_experts, bias=True)
 
     def get_periodogram(self, inputs, ker_len=50, con=1, n=10000):
-        if inputs.dim() == 2:
-            x_0 = inputs.unsqueeze(2)
-        else:
-            x_0 = inputs
-        x_0 = x_0 - torch.mean(x_0, dim=1, keepdim=True)
-
-        v = torch.arange(0, n) / n
-        if con:
-            if ker_len is None:
-                ker_len = n // 4
-                ker_len = min(ker_len, 50)
-
-            x_0 = x_0.permute(0, 2, 1)
-            ker = (torch.ones(1, 1, ker_len) / ker_len).to(x_0.device)
-            x_c = F.conv1d(x_0, ker, padding="same")
-            x_c[:, :, :ker_len // 2] = x_c[:, :, ker_len // 2:ker_len // 2 + 1]
-            x_c[:, :, -ker_len // 2:] = x_c[:, :, -ker_len // 2 - 1:-ker_len // 2]
-            x_0 = x_0 - x_c
-            x_0 = x_0.permute(0, 2, 1)
-
-        dft = torch.fft.fft(x_0, dim=1, n=n) / np.sqrt(n)
-        dft = dft[:, :n//2, :]
-        I = torch.abs(dft) ** 2
-
-        I_sum = torch.sum(I, dim=1, keepdim=True)
-        I_sum[I_sum == 0] = 1
-        I = I / I_sum
-
-        if torch.any(I_sum == 0):
-            print("Zeros in the sum")
-            raise ValueError
-
-        if inputs.dim() == 2:
+        n_fft   = 128
+        ker_len =12
+        if inputs.ndim == 2:          # (B, L)  → (B, L, 1)
+            x = inputs.unsqueeze(2)
+        else:                         # already (B, L, C)
+            x = inputs
+    
+        B, L, C = x.shape
+        x = x - x.mean(dim=1, keepdim=True)
+    
+        # ------------------------------------------------------------------ parameters
+        if n_fft is None:
+            n_fft = 1 << (L - 1).bit_length()          # next power-of-two ≥ L
+        if ker_len is None:
+            ker_len = min(L // 4, 50)                  # never larger than the signal
+        ker_half = ker_len // 2
+    
+        # ------------------------------------------------------------------ detrend
+        if con and ker_len > 0:
+            #  (B, L, C) → (B, C, L) for conv1d
+            x_perm = x.permute(0, 2, 1)
+            ker    = torch.ones(1, 1, ker_len, device=x.device) / ker_len
+            trend  = F.conv1d(x_perm, ker, padding="same")
+            # Clamp boundary copies so we don’t index out of range for short signals
+            left  = min(ker_half, L - 1)
+            right = min(ker_half, L - 1)
+            trend[:, :, :left]      = trend[:, :, left:left+1]
+            trend[:, :, -right:]    = trend[:, :, -(right+1):-right]
+            x_detrended = x_perm - trend
+            x = x_detrended.permute(0, 2, 1)           # back to (B, L, C)
+    
+        # ------------------------------------------------------------------ FFT
+        dft  = torch.fft.fft(x, n=n_fft, dim=1) / np.sqrt(n_fft)
+        dft  = dft[:, : n_fft // 2, :]                 # keep positive freqs
+        I    = torch.abs(dft) ** 2                    # periodogram
+    
+        # ------------------------------------------------------------------ normalise
+        I_sum = I.sum(dim=1, keepdim=True)
+        I_sum[I_sum == 0] = 1                         # avoid /0
+        I /= I_sum
+    
+        # ------------------------------------------------------------------ squeeze back if original was 2-D
+        if inputs.ndim == 2:
             I = I.squeeze(2)
-            
+    
         return I