Update modeling_super_linear.py

modeling_super_linear.py

@@ -295,50 +295,62 @@ class SparseNoisyMoE(nn.Module):
     def get_periodogram(self, inputs, ker_len=50, con=1, n=10000):
         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
+        if inputs.ndim == 2:                   # (B, L)
+            B, L          = inputs.shape
+            C             = 1
+            x             = inputs.unsqueeze(2)            # → (B, L, 1)
+            time_first    = True                           # time is dim-1
+        elif inputs.ndim == 3:
+            B, d1, d2     = inputs.shape
+            if d1 < d2:                                    # (B, L, C)
+                L, C  = d1, d2
+                x      = inputs
+                time_first = True
+            else:                                          # (B, C, L)
+                C, L  = d1, d2
+                x      = inputs.transpose(1, 2)            # → (B, L, C)
+                time_first = False
+        else:
+            raise ValueError("Input must be (B,L), (B,L,C) or (B,C,L)")
     
-        B, L, C = x.shape
+        # ---------- centre the signal ----------
         x = x - x.mean(dim=1, keepdim=True)
     
-        # ------------------------------------------------------------------ parameters
+        # ---------- parameter defaults ----------
         if n_fft is None:
-            n_fft = 1 << (L - 1).bit_length()          # next power-of-two ≥ L
+            n_fft = 1 << (L - 1).bit_length()
         if ker_len is None:
-            ker_len = min(L // 4, 50)                  # never larger than the signal
+            ker_len = min(L // 4, 50)
         ker_half = ker_len // 2
     
-        # ------------------------------------------------------------------ detrend
+        # ---------- high-pass detrend ----------
         if con and ker_len > 0:
-            #  (B, L, C) → (B, C, L) for conv1d
-            x_perm = x.permute(0, 2, 1)
+            x_perm = x.permute(0, 2, 1)              # (B, C, L)
             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)
+            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      = (x_perm - trend).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
+        # ---------- FFT ----------
+        dft = torch.fft.fft(x, n=n_fft, dim=1) / np.sqrt(n_fft)
+        I   = (dft[:, : n_fft//2, :]).abs() ** 2
     
-        # ------------------------------------------------------------------ normalise
+        # ---------- normalise ----------
         I_sum = I.sum(dim=1, keepdim=True)
-        I_sum[I_sum == 0] = 1                         # avoid /0
+        I_sum[I_sum == 0] = 1
         I /= I_sum
     
-        # ------------------------------------------------------------------ squeeze back if original was 2-D
-        if inputs.ndim == 2:
-            I = I.squeeze(2)
+        # ---------- restore original layout ----------
+        if inputs.ndim == 2:            # wanted (B, … )
+            return I.squeeze(2)
     
-        return I
+        if time_first:                  # original was (B, L, C)
+            return I                    # already (B, F, C)
+        else:                           # original was (B, C, L) → (B, C, F)
+            return I.transpose(1, 2)
 
 
     def fourier_interp_dim1(self,x, target_len: int = 512):