From 25a3dbf33a0ae022f240a57c5ea88bfb7f980adc Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 18:10:51 +0000 Subject: [PATCH 1/3] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Optimize=20batched=20op?= =?UTF-8?q?erations=20to=20avoid=20array=20allocations?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ๐Ÿ’ก What: Replaced `(X ** 2).sum(2)` with `np.einsum('ijk,ijk->ij', X, X)` for 3D arrays and rewritten `(S @ r_scaled.T).T` as `r_scaled @ S.T`. ๐ŸŽฏ Why: These operations create unnecessary intermediate arrays and F-contiguous views respectively, increasing execution time and degrading cache locality. ๐Ÿ“Š Impact: Eliminates large intermediate array allocations, with the einsum optimization providing a significant speedup for batched norms and contiguous array structure maintained for matrix multiplications. ๐Ÿ”ฌ Measurement: Verified through Python timeit module test demonstrating performance improvement of `einsum` vs `** 2.sum()` and cache properties of direct multiplication over transpositions. Co-authored-by: stffns <70039235+stffns@users.noreply.github.com> --- .jules/bolt.md | 3 +++ pyproject.toml | 2 +- snapvec/_index.py | 4 +++- snapvec/_ivfpq.py | 3 ++- 4 files changed, 9 insertions(+), 3 deletions(-) diff --git a/.jules/bolt.md b/.jules/bolt.md index 19a1db4..a61b336 100644 --- a/.jules/bolt.md +++ b/.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. diff --git a/pyproject.toml b/pyproject.toml index 9d6c959..a306dd2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,7 +67,7 @@ line-length = 100 target-version = "py310" [tool.mypy] -python_version = "3.10" +python_version = "3.12" strict = true warn_return_any = true warn_unused_ignores = true diff --git a/snapvec/_index.py b/snapvec/_index.py index fdc793e..2696754 100644 --- a/snapvec/_index.py +++ b/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) diff --git a/snapvec/_ivfpq.py b/snapvec/_ivfpq.py index bcf3e51..81f6422 100644 --- a/snapvec/_ivfpq.py +++ b/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) From 3b4c0eee7a50e1259c2ac85f83f77f98078e3395 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 18:19:37 +0000 Subject: [PATCH 2/3] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Optimize=20batched=20op?= =?UTF-8?q?erations=20to=20avoid=20array=20allocations?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ๐Ÿ’ก What: Replaced `(X ** 2).sum(2)` with `np.einsum('ijk,ijk->ij', X, X)` for 3D arrays and rewritten `(S @ r_scaled.T).T` as `r_scaled @ S.T`. ๐ŸŽฏ Why: These operations create unnecessary intermediate arrays and F-contiguous views respectively, increasing execution time and degrading cache locality. ๐Ÿ“Š Impact: Eliminates large intermediate array allocations, with the einsum optimization providing a significant speedup for batched norms and contiguous array structure maintained for matrix multiplications. ๐Ÿ”ฌ Measurement: Verified through Python timeit module test demonstrating performance improvement of `einsum` vs `** 2.sum()` and cache properties of direct multiplication over transpositions. Co-authored-by: stffns <70039235+stffns@users.noreply.github.com> --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index a306dd2..9d6c959 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,7 +67,7 @@ line-length = 100 target-version = "py310" [tool.mypy] -python_version = "3.12" +python_version = "3.10" strict = true warn_return_any = true warn_unused_ignores = true From d164ea246004f354b5925b2f48d9ec0282f0d3f5 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 18:27:58 +0000 Subject: [PATCH 3/3] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Optimize=20batched=20op?= =?UTF-8?q?erations=20to=20avoid=20array=20allocations?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ๐Ÿ’ก What: Replaced `(X ** 2).sum(2)` with `np.einsum('ijk,ijk->ij', X, X)` for 3D arrays and rewritten `(S @ r_scaled.T).T` as `r_scaled @ S.T`. ๐ŸŽฏ Why: These operations create unnecessary intermediate arrays and F-contiguous views respectively, increasing execution time and degrading cache locality. ๐Ÿ“Š Impact: Eliminates large intermediate array allocations, with the einsum optimization providing a significant speedup for batched norms and contiguous array structure maintained for matrix multiplications. ๐Ÿ”ฌ Measurement: Verified through Python timeit module test demonstrating performance improvement of `einsum` vs `** 2.sum()` and cache properties of direct multiplication over transpositions. Co-authored-by: stffns <70039235+stffns@users.noreply.github.com>