⚡ Bolt: Optimize batched operations to avoid array allocations

.jules/bolt.md

@@ -1,3 +1,6 @@
 ## 2024-05-18 - Fast row-wise Euclidean norm in pure NumPy
 **Learning:** In performance-critical paths, computing the batch norm of a 2D array via `np.linalg.norm(arr, axis=1)` is relatively slow. Using `np.sqrt(np.einsum('ij,ij->i', arr, arr))` is significantly faster (~4x speedup on a laptop CPU for typical batch sizes). If `keepdims=True` behavior is needed, appending `[:, np.newaxis]` matches the original shape seamlessly.
 **Action:** Always prefer `np.sqrt(np.einsum('ij,ij->i', arr, arr))` over `np.linalg.norm(arr, axis=1)` when computing row-wise vector norms in NumPy to eliminate dispatch overhead and improve execution speed.
+## 2025-02-18 - Fast batched squared norms for 3D arrays and Contiguous Matrix Multiplications
+**Learning:** Using `np.einsum('ijk,ijk->ij', X, X)` is significantly faster than `(X ** 2).sum(2)` for calculating squared Euclidean norms along the last axis of a 3D NumPy array, avoiding large intermediate array allocations. Also, explicitly using associativity (e.g. `R @ S.T` instead of `(S @ R.T).T`) yields a C-contiguous array instead of an F-contiguous view, which speeds up the operation and subsequent steps relying on cache locality.
+**Action:** Replace `(X ** 2).sum(2)` with `np.einsum('ijk,ijk->ij', X, X)` for 3D array batched squared norm calculations and rewrite expression to avoid explicit transpositions to preserve C-contiguity in critical data paths.

snapvec/_index.py

@@ -303,7 +303,9 @@ def add_batch(self, ids: list[Any], vectors: NDArray[np.float32]) -> None:
             reconstructed: NDArray[np.float32] = self._centroids[batch_idx]
             r_scaled: NDArray[np.float32] = scaled - reconstructed
             # sign(S·r_rot) = sign(S·r_scaled) — scale-invariant
-            S_r: NDArray[np.float32] = (self._S @ r_scaled.T).T
+            # Optimized: Avoid explicit transpositions and intermediate allocations.
+            # Yields a C-contiguous array rather than an F-contiguous view.
+            S_r: NDArray[np.float32] = r_scaled @ self._S.T
             qjl_signs = np.sign(S_r).astype(np.int8)
             qjl_signs[qjl_signs == 0] = 1
             # Store ‖r_rot‖ = ‖r_scaled‖/√pdim (unscaled space norm)

snapvec/_ivfpq.py

@@ -429,7 +429,8 @@ def add_batch(
             if self.keep_full_precision else
             np.empty((0, self._pdim), dtype=np.float16)
         )
-        cb_norms = (self._codebooks ** 2).sum(2)             # (M, K)
+        # Optimized: np.einsum('ijk,ijk->ij', X, X) is faster and avoids large intermediate arrays
+        cb_norms = np.einsum('ijk,ijk->ij', self._codebooks, self._codebooks)             # (M, K)
         cb_T = np.transpose(self._codebooks, (0, 2, 1))      # (M, d_sub, K)
         for start in range(0, n, self._ENCODE_CHUNK):
             end = min(start + self._ENCODE_CHUNK, n)