⚡ Bolt: Optimize euclidean distance loop using np.einsum

.jules/bolt.md

@@ -48,3 +48,6 @@
 ## 2025-05-19 - Dot product scalar reductions in MMLE M-step
 **Learning:** During GPCM M-step item gradient and expected log-likelihood calculations, `float(np.sum(r_counts * lp))` and `float(np.sum((resid @ scores) * base))` construct full intermediate arrays of shape `(N, K)` and `(N,)` respectively before reducing them to a scalar sum.
 **Action:** Replace `np.sum(A * B)` with `np.vdot(A, B)` when calculating a scalar reduction over an element-wise product of arrays with identical shapes. This entirely skips allocating the intermediate product array and improves M-step computation speeds significantly.
+## 2025-05-19 - Fast sum of squares with einsum
+**Learning:** Computing the sum of squares across an axis using `np.sum(diff * diff, axis=1)` allocates an intermediate 2D array, which becomes a bottleneck in gradient/distance loops.
+**Action:** Replace `np.sum(diff * diff, axis=1)` with `np.einsum('ij,ij->i', diff, diff)` to compute the squared differences across the axis without allocating an intermediate `(N, K)` array, improving execution speed and reducing memory bandwidth.

python/fast_mlsirm/estimators/marginal.py

@@ -1042,7 +1042,8 @@ def eta_of(alpha_c: float, b_c: float, zeta_c: np.ndarray) -> np.ndarray:
                         deta_z = x_grid  # (Nx, K)
                     else:
                         diff = x_grid - zeta_i[None, :]
-                        dist = np.sqrt(eps_distance + np.sum(diff * diff, axis=1))
+                        # Optimized distance calculation: replace memory allocation in np.sum(diff * diff, axis=1) with np.einsum
+                        dist = np.sqrt(eps_distance + np.einsum('ij,ij->i', diff, diff))
                         deta_z = gamma * diff / dist[:, None]  # (Nx, K)
                     g_zeta = (
                         np.einsum("stx,xk->k", resid, deta_z, optimize=True)