fixing welford combine for spectra calculation

Author: azrael417

data_process/data_process_helpers.py

@@ -132,9 +132,13 @@ def welford_combine(stats1, stats2):
         if (s_b["counts"].ndim != 0) and (s_a["counts"].ndim != s_a["values"].ndim):
             n_a = s_a["counts"][None, :, None, None]
             n_b = s_b["counts"][None, :, None, None]
+            reshape = True
         else:
             n_a = s_a["counts"]
             n_b = s_b["counts"]
+            reshape = False
+
+        # combined counts
         n_ab = n_a + n_b
 
         if s_a["type"] == "min":
@@ -157,7 +161,10 @@ def welford_combine(stats1, stats2):
                 ], dim=0
             ).contiguous()
 
-        stats[k] = {"counts": n_ab.reshape(-1),
+        if reshape:
+            n_ab = n_ab.reshape(-1)
+
+        stats[k] = {"counts": n_ab,
                     "type": s_a["type"],
                     "values": values}
 

data_process/get_spectra.py

@@ -36,10 +36,11 @@
 from wb2_helpers import DistributedProgressBar
 from data_process_helpers import welford_combine, collective_reduce, binary_reduce
 
+@torch.compile(fullgraph=True)
 def compute_powerspectrum(x, sht):
-    coeffs = sht(x).abs().pow(2)
+    coeffs = torch.square(torch.abs(sht(x)))
     coeffs[..., 1:] *= 2.0
-    power_spectrum = coeffs.sum(dim=-1)
+    power_spectrum = torch.sum(coeffs, dim=-1)
 
     return power_spectrum