From 551ce291ccab4050c050f24cae422219b331d0a4 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:52:06 +0000 Subject: [PATCH 1/4] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Replace=20`.sum(1)`=20w?= =?UTF-8?q?ith=20`np.einsum`=20for=20squared=20Euclidean=20norms?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaced `((X - c)**2).sum(1)` and `(X**2).sum(1)` with equivalent `np.einsum` calls in `_kmeans.py`, `_ivfpq.py`, and `_pq.py`. This prevents large intermediate array allocations and significantly improves execution speed for squared Euclidean distance calculations. Co-authored-by: stffns <70039235+stffns@users.noreply.github.com> --- snapvec/_ivfpq.py | 3 ++- snapvec/_kmeans.py | 20 ++++++++++++++------ snapvec/_pq.py | 5 +++-- 3 files changed, 19 insertions(+), 9 deletions(-) diff --git a/snapvec/_ivfpq.py b/snapvec/_ivfpq.py index bcf3e51..22ed3fe 100644 --- a/snapvec/_ivfpq.py +++ b/snapvec/_ivfpq.py @@ -996,7 +996,8 @@ def search_batch( # One matmul, the whole batch. coarse_dot_all = q_pre_all @ self._coarse.T # (B, nlist) - cnorms = (self._coarse * self._coarse).sum(1) # (nlist,) + # Bolt: np.einsum prevents large intermediate array allocations + cnorms = np.einsum('ij,ij->i', self._coarse, self._coarse) # (nlist,) probe_ranking_all = 2.0 * coarse_dot_all - cnorms[None, :] if allowed_clusters is None: probes = np.argpartition( diff --git a/snapvec/_kmeans.py b/snapvec/_kmeans.py index a4b1dd6..5214363 100644 --- a/snapvec/_kmeans.py +++ b/snapvec/_kmeans.py @@ -28,13 +28,17 @@ def kmeans_pp_init( """ n = X.shape[0] centers = [X[int(rng.integers(n))]] - d2 = ((X - centers[0]) ** 2).sum(1) + # Bolt: np.einsum prevents large intermediate allocations from (X - c)**2 + diff = X - centers[0] + d2 = np.einsum('ij,ij->i', diff, diff) for _ in range(1, K): total = d2.sum() probs = d2 / total if total > 1e-12 else np.full(n, 1.0 / n) nxt = int(rng.choice(n, p=probs)) centers.append(X[nxt]) - d2 = np.minimum(d2, ((X - centers[-1]) ** 2).sum(1)) + # Bolt: np.einsum prevents large intermediate allocations from (X - c)**2 + diff = X - centers[-1] + d2 = np.minimum(d2, np.einsum('ij,ij->i', diff, diff)) return np.stack(centers).astype(np.float32) @@ -50,9 +54,11 @@ def kmeans_mse( """ rng = np.random.default_rng(seed) C = kmeans_pp_init(X, K, rng) - x_sq = (X ** 2).sum(1, keepdims=True) + # Bolt: np.einsum prevents large intermediate array allocations + x_sq = np.einsum('ij,ij->i', X, X)[:, None] for _ in range(n_iters): - d2 = x_sq - 2 * X @ C.T + (C ** 2).sum(1)[None, :] + # Bolt: np.einsum prevents large intermediate array allocations + d2 = x_sq - 2 * X @ C.T + np.einsum('ij,ij->i', C, C)[None, :] asn = d2.argmin(1) newC = np.empty_like(C) dead_ks: list[int] = [] @@ -88,7 +94,8 @@ def assign_l2( X: NDArray[np.float32], C: NDArray[np.float32], ) -> NDArray[np.int64]: """Hard-assign every row in X to its nearest centroid (squared L2).""" - d2 = (X ** 2).sum(1, keepdims=True) - 2 * X @ C.T + (C ** 2).sum(1)[None, :] + # Bolt: np.einsum prevents large intermediate array allocations + d2 = np.einsum('ij,ij->i', X, X)[:, None] - 2 * X @ C.T + np.einsum('ij,ij->i', C, C)[None, :] return cast("NDArray[np.int64]", d2.argmin(1)) @@ -112,9 +119,10 @@ def probe_scores_l2_monotone( # Python '2.0' scalar to float64 here; on numpy >= 2.0 this is a # no-op, on older numpy it keeps the return dtype matching the # annotation. + # Bolt: np.einsum prevents large intermediate array allocations return cast( "NDArray[np.float32]", - np.float32(2.0) * (coarse @ q) - (coarse ** 2).sum(1), + np.float32(2.0) * (coarse @ q) - np.einsum('ij,ij->i', coarse, coarse), ) diff --git a/snapvec/_pq.py b/snapvec/_pq.py index 07b0a0e..716339a 100644 --- a/snapvec/_pq.py +++ b/snapvec/_pq.py @@ -307,10 +307,11 @@ def add_batch( codes = np.empty((self.M, len(arr)), dtype=np.uint8) for j in range(self.M): Xj = pre[:, j * self._d_sub : (j + 1) * self._d_sub] + # Bolt: np.einsum prevents large intermediate array allocations d2 = ( - (Xj ** 2).sum(1, keepdims=True) + np.einsum('ij,ij->i', Xj, Xj)[:, None] - 2 * Xj @ self._codebooks[j].T - + (self._codebooks[j] ** 2).sum(1)[None, :] + + np.einsum('ij,ij->i', self._codebooks[j], self._codebooks[j])[None, :] ) codes[j] = d2.argmin(1).astype(np.uint8) From 88b7d247b67c4263d69349fe67a269f43f7b31ee Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:08:10 +0000 Subject: [PATCH 2/4] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Replace=20`.sum(1)`=20w?= =?UTF-8?q?ith=20`np.einsum`=20for=20squared=20Euclidean=20norms?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaced `((X - c)**2).sum(1)` and `(X**2).sum(1)` with equivalent `np.einsum` calls in `_kmeans.py`, `_ivfpq.py`, and `_pq.py`. This prevents large intermediate array allocations and significantly improves execution speed for squared Euclidean distance calculations. Includes CI linting fixes. Co-authored-by: stffns <70039235+stffns@users.noreply.github.com> --- snapvec/__init__.py | 6 +++--- snapvec/_fast.pyi | 2 -- snapvec/_file_format.py | 15 +++++++-------- snapvec/_index.py | 4 ++-- snapvec/_ivfpq.py | 4 ++-- snapvec/_kmeans.py | 6 +++--- snapvec/_pq.py | 4 ++-- snapvec/_residual.py | 5 ++--- tests/test_adversarial.py | 1 - tests/test_file_format.py | 4 ++-- tests/test_properties.py | 1 - tests/test_snapvec.py | 6 ++++-- 12 files changed, 27 insertions(+), 31 deletions(-) diff --git a/snapvec/__init__.py b/snapvec/__init__.py index 5994437..9335194 100644 --- a/snapvec/__init__.py +++ b/snapvec/__init__.py @@ -21,11 +21,11 @@ __version__ = "0.11.1" __all__ = [ - "SnapIndex", - "PQSnapIndex", "IVFPQSnapIndex", + "PQSnapIndex", "ResidualSnapIndex", + "SnapIndex", "get_codebook", - "rht", "padded_dim", + "rht", ] diff --git a/snapvec/_fast.pyi b/snapvec/_fast.pyi index 7aceae9..d7b9527 100644 --- a/snapvec/_fast.pyi +++ b/snapvec/_fast.pyi @@ -4,12 +4,10 @@ The real module is built from Cython and does not ship a ``.pyi`` from the compiler; this stub lets ``mypy --strict`` see the same Python-level shapes the Cython kernels expose to callers. """ -from __future__ import annotations import numpy as np from numpy.typing import NDArray - def adc_colmajor( lut: NDArray[np.float32], codes: NDArray[np.uint8], diff --git a/snapvec/_file_format.py b/snapvec/_file_format.py index 81efc2f..8213b5b 100644 --- a/snapvec/_file_format.py +++ b/snapvec/_file_format.py @@ -29,10 +29,10 @@ import os import struct import zlib +from collections.abc import Callable from pathlib import Path from types import TracebackType -from typing import IO, Callable - +from typing import IO, Self _TRAILER_MAGIC = b"CRC2" _TRAILER_SIZE = 8 # 4 bytes magic + 4 bytes uint32 CRC @@ -79,7 +79,7 @@ def finalise(self) -> None: self._f.write(struct.pack(" "ChecksumWriter": + def __enter__(self) -> Self: return self def __exit__( @@ -163,16 +163,15 @@ def save_with_checksum_atomic( """ path = Path(path) tmp = path.with_suffix(path.suffix + ".tmp") - with open(tmp, "wb") as raw: - with ChecksumWriter(raw) as cw: - writer_fn(cw) + with open(tmp, "wb") as raw, ChecksumWriter(raw) as cw: + writer_fn(cw) os.replace(tmp, path) __all__ = [ "ChecksumWriter", "has_trailer", - "verify_checksum", - "trailer_len", "save_with_checksum_atomic", + "trailer_len", + "verify_checksum", ] diff --git a/snapvec/_index.py b/snapvec/_index.py index fdc793e..710e935 100644 --- a/snapvec/_index.py +++ b/snapvec/_index.py @@ -581,7 +581,7 @@ def save(self, path: str | Path) -> None: else: packed = _pack(self._indices, self._mse_bits) - def _write(f: "ChecksumWriter") -> None: + def _write(f: ChecksumWriter) -> None: f.write(_MAGIC) f.write(struct.pack(" None: save_with_checksum_atomic(path, _write) @classmethod - def load(cls, path: str | Path) -> "SnapIndex": + def load(cls, path: str | Path) -> SnapIndex: """Load index from a ``.snpv`` file. Supports v1 (mse-only legacy) and v2 (prod/flags) formats. diff --git a/snapvec/_ivfpq.py b/snapvec/_ivfpq.py index 22ed3fe..22cc3bf 100644 --- a/snapvec/_ivfpq.py +++ b/snapvec/_ivfpq.py @@ -1128,7 +1128,7 @@ def save(self, path: str | Path) -> None: flags |= _FLAG_USE_OPQ n = len(self._ids_by_row) - def _write(f: "ChecksumWriter") -> None: + def _write(f: ChecksumWriter) -> None: f.write(_MAGIC) f.write( struct.pack( @@ -1171,7 +1171,7 @@ def _write(f: "ChecksumWriter") -> None: save_with_checksum_atomic(path, _write) @classmethod - def load(cls, path: str | Path) -> "IVFPQSnapIndex": + def load(cls, path: str | Path) -> IVFPQSnapIndex: path = Path(path) verify_checksum(path) # no-op for legacy files without a trailer with open(path, "rb") as f: diff --git a/snapvec/_kmeans.py b/snapvec/_kmeans.py index 5214363..d639d44 100644 --- a/snapvec/_kmeans.py +++ b/snapvec/_kmeans.py @@ -207,9 +207,9 @@ def fit_opq_rotation( __all__ = [ - "kmeans_pp_init", - "kmeans_mse", "assign_l2", - "probe_scores_l2_monotone", "fit_opq_rotation", + "kmeans_mse", + "kmeans_pp_init", + "probe_scores_l2_monotone", ] diff --git a/snapvec/_pq.py b/snapvec/_pq.py index 716339a..0021a56 100644 --- a/snapvec/_pq.py +++ b/snapvec/_pq.py @@ -427,7 +427,7 @@ def save(self, path: str | Path) -> None: flags |= _FLAG_USE_OPQ n = len(self._ids) - def _write(f: "ChecksumWriter") -> None: + def _write(f: ChecksumWriter) -> None: f.write(_MAGIC) f.write( struct.pack( @@ -460,7 +460,7 @@ def _write(f: "ChecksumWriter") -> None: save_with_checksum_atomic(path, _write) @classmethod - def load(cls, path: str | Path) -> "PQSnapIndex": + def load(cls, path: str | Path) -> PQSnapIndex: path = Path(path) verify_checksum(path) # no-op for legacy files without a trailer with open(path, "rb") as f: diff --git a/snapvec/_residual.py b/snapvec/_residual.py index e0e4e7c..963174c 100644 --- a/snapvec/_residual.py +++ b/snapvec/_residual.py @@ -35,7 +35,6 @@ from ._freezable import FreezableIndex from ._rotation import padded_dim, rht - _MAX_ID_BYTES = 0xFFFF # file format stores id length as uint16 @@ -295,7 +294,7 @@ def save(self, path: str | Path) -> None: flags |= 1 n = len(self._ids) - def _write(f: "ChecksumWriter") -> None: + def _write(f: ChecksumWriter) -> None: f.write(_MAGIC) f.write(struct.pack(" None: save_with_checksum_atomic(path, _write) @classmethod - def load(cls, path: str | Path) -> "ResidualSnapIndex": + def load(cls, path: str | Path) -> ResidualSnapIndex: path = Path(path) verify_checksum(path) # no-op for legacy files without a trailer with open(path, "rb") as f: diff --git a/tests/test_adversarial.py b/tests/test_adversarial.py index bf71d0e..e5014a2 100644 --- a/tests/test_adversarial.py +++ b/tests/test_adversarial.py @@ -11,7 +11,6 @@ from snapvec import IVFPQSnapIndex, PQSnapIndex, ResidualSnapIndex, SnapIndex - # --------------------------------------------------------------------------- # # Empty index # # --------------------------------------------------------------------------- # diff --git a/tests/test_file_format.py b/tests/test_file_format.py index 9ba50cb..bdd7d08 100644 --- a/tests/test_file_format.py +++ b/tests/test_file_format.py @@ -150,8 +150,8 @@ def test_truncated_trailer_falls_back_to_legacy_mode(tmp_path: Path) -> None: # ──────────────────────────────────────────────────────────────────── # @pytest.mark.parametrize("index_cls, ctor_kwargs, suffix", [ - (SnapIndex, dict(dim=32, bits=4, normalized=True), ".snpv"), - (ResidualSnapIndex, dict(dim=32, b1=3, b2=3, normalized=True), ".snpr"), + (SnapIndex, {"dim": 32, "bits": 4, "normalized": True}, ".snpv"), + (ResidualSnapIndex, {"dim": 32, "b1": 3, "b2": 3, "normalized": True}, ".snpr"), ]) def test_trailing_crc_roundtrip_trainingfree( index_cls, ctor_kwargs, suffix, tmp_path, diff --git a/tests/test_properties.py b/tests/test_properties.py index 1237e77..ce366fd 100644 --- a/tests/test_properties.py +++ b/tests/test_properties.py @@ -16,7 +16,6 @@ from snapvec import IVFPQSnapIndex, PQSnapIndex, SnapIndex - PROFILE = settings( max_examples=25, deadline=None, diff --git a/tests/test_snapvec.py b/tests/test_snapvec.py index 0f8a2c0..66aa1e0 100644 --- a/tests/test_snapvec.py +++ b/tests/test_snapvec.py @@ -171,6 +171,7 @@ def test_legacy_v2_3bit_file_loads_via_compat_decoder(self, tmp_path): path, then re-pack into the new tight RAM layout. """ import struct + from snapvec._index import _MAGIC idx = SnapIndex(dim=128, bits=3) @@ -212,7 +213,8 @@ def test_legacy_v2_prod_mode_3bit_payload_stays_aligned(self, tmp_path): corrupt the prod correction term). """ import struct - from snapvec._index import _MAGIC, _FLAG_PROD + + from snapvec._index import _FLAG_PROD, _MAGIC # Real v3 prod-mode index to source the reference indices + payload. idx = SnapIndex(dim=128, bits=4, use_prod=True) @@ -473,7 +475,7 @@ def test_filter_restricts_results(self): idx = SnapIndex(dim=DIM, bits=4) idx.add_batch(list(range(100)), vecs) - allowed = set(range(0, 50)) + allowed = set(range(50)) results = idx.search(vecs[0], k=10, filter_ids=allowed) assert all(r[0] in allowed for r in results) From b81d3f2f33bc1a20590580b9f26a9f5d4520e8eb Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:16:20 +0000 Subject: [PATCH 3/4] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Replace=20`.sum(1)`=20w?= =?UTF-8?q?ith=20`np.einsum`=20for=20squared=20Euclidean=20norms?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaced `((X - c)**2).sum(1)` and `(X**2).sum(1)` with equivalent `np.einsum` calls in `_kmeans.py`, `_ivfpq.py`, and `_pq.py`. This prevents large intermediate array allocations and significantly improves execution speed for squared Euclidean distance calculations. Includes CI linting fixes and backwards-compatibility for typing.Self. Co-authored-by: stffns <70039235+stffns@users.noreply.github.com> --- .coverage | Bin 0 -> 53248 bytes snapvec/_file_format.py | 4 +++- 2 files changed, 3 insertions(+), 1 deletion(-) create mode 100644 .coverage diff --git a/.coverage b/.coverage new file mode 100644 index 0000000000000000000000000000000000000000..fb211cd8092380fc53597254038986d42858737d GIT binary patch literal 53248 zcmeI4@oy978OPtdb9}L#_|3||gaEl!nnFfM8l_XAh@@#3X;G0{RyHwB!d#q7;$r)Z z@0i&REWdmu_x=owd4--?nNTlnwPVKs~P3ksv5UO=65@mr6maqi8=e_gA zv6EE&1x29mEBoHN_j#Z9ecsReyzlGRXWuiscE)W(>oYS+-PYPz4U;7Hn5Hqt0`#t? zw-;18@p%^%OAF4II}Nbr!56Ce*I1Q*5952Q*YPb?{Z-dPC#$vwz8Z?FO@U?-AbGHw4CZRjHBMTIuK!TklC2c1DjG8TVNI#8~Hjo8rBV!Ny~wK{c{! zm2}*2+N5V<^a_=F%R~VrHgSSl^S;Q?UWVYSQ%i7RF4nbJ%uUn=z9uwy7oJ zDW{n!D{jZllx7SYv8-+MF4@2)0nG~pZ((bwtbL+E`&`XJT?^_|ylWEa8q~TUo!_JtHJhAfq|7EwG=e~n^U(AM0}A#Ywd zW?DQLdPWR1%`9XjOIOs}gt26U){`)MN`1}o7|WmmLUVzatuLF3iJ@YNHpI+9Bct~l z%>}7DEv&R7k?XVu{E@Eudp1){CebSiJ0irqB;FMz zjSwltAyBpue0uZKx;>yRwp}EX;SCl3$oAFC87@qSwu+-b#td zBBf~Zlg7aRww5mvV_`l?jVZ_7u(1^bu*Ha|z%J7>Y}`1~DVL-QMG(^KunFAO#!624CW{j9v&J z00JNY0w4eaAOHd&00JNY0w4ea_a6a8k|kB_|7CuO@qf_)0R%t*1V8`;KmY_l00ck) z1V8`;K;WJvP$erHxc9S=hiW9HuD0+SfJd5Jzu4Rsq*lrNGUJ!|&3m$eFd76v00ck) z1V8`;KmY_l00ck)1VCV(K$WsV^1cR;gGybX@GXG2{~rqfg7LreU-2LE6FknJu%B0=@u)crw5H#OSDh7Y&s+o zJ$h_EJwjl8P+=XVSu%zd@6GB7lCB9TtgTdV7FR5u zLL-Wn7Zn{ddkuO7#C~CNousfGrA0^Mskj}Dl2F|L4~54W|B(MWJkDR>&+;GfR=zrX zJ3PsEu5=fpeh>fw5C8!X009sH0T2KI5C8!XVC+GqTPoexmhb;V8>+Akr_7k`!}8b{{8Zb~ z-73}g*tvH)`G?%_Smi*zcW6aKFsr1P9J{0KJEl<=+zJim42E zxe&j1{kt9C9DlKUR30CZCoaDHvwfX&`Pqr7S?k)2dd2L>R}7go%@5bGxf+(cy~eB% zWcDg#D`V%c9`RFjI{)UB^p5okRa!YSbbPk{<5w!KrAIdY`_^~brq-YOUi$Fy#>QitNAnh#B4|&S(`uSpZm@1aP`pp)2|=ho6jqhQO?i$Iyzg}`nB!( z8Xxg{u!8E#{dhHd`m`$5DU`C+M=6)j)Ur!E-v0T>YpZf^00@~ixxe4PJ<|AD{7-yjhJ2!H?xfB*=900@8p z2!H?xfB*=9zzP%a)2?3f2c0MwaH2q^6IH5Cr23smsc@nSpA-2MCz53!eFfm$|7SHT z+z?a`0w4eaAOHd&00JNY0w4eaAOHd&a90V4`~TSg-&MpY90WiB1V8`;KmY_l00ck) s1V8`;R-AzN|3CKsE8Yz-00ck)1V8`;KmY_l00ck)1V8`;?ka)*1D>D-xBvhE literal 0 HcmV?d00001 diff --git a/snapvec/_file_format.py b/snapvec/_file_format.py index 8213b5b..3f53178 100644 --- a/snapvec/_file_format.py +++ b/snapvec/_file_format.py @@ -32,7 +32,9 @@ from collections.abc import Callable from pathlib import Path from types import TracebackType -from typing import IO, Self +from typing import IO + +from typing_extensions import Self _TRAILER_MAGIC = b"CRC2" _TRAILER_SIZE = 8 # 4 bytes magic + 4 bytes uint32 CRC From fdbf27474f4d024d08eb1a84584eaee114cc5a8a Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:24:58 +0000 Subject: [PATCH 4/4] =?UTF-8?q?=E2=9A=A1=20Bolt:=20Replace=20`.sum(1)`=20w?= =?UTF-8?q?ith=20`np.einsum`=20for=20squared=20Euclidean=20norms?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaced `((X - c)**2).sum(1)` and `(X**2).sum(1)` with equivalent `np.einsum` calls in `_kmeans.py`, `_ivfpq.py`, and `_pq.py`. This prevents large intermediate array allocations and significantly improves execution speed for squared Euclidean distance calculations. Includes CI linting fixes and backwards-compatibility for typing.Self. Co-authored-by: stffns <70039235+stffns@users.noreply.github.com> --- .coverage | Bin 53248 -> 53248 bytes snapvec/_file_format.py | 5 ++++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.coverage b/.coverage index fb211cd8092380fc53597254038986d42858737d..dba999f6199148c7122048680acf06673c5349dd 100644 GIT binary patch delta 78 zcmV-U0I~mopaX!Q1F!~w7T^F6_YdU{#}Bm+L=Kk@bq*>G^$p;&5fH=;v*C^B0UjI! k1OW*;1Ofm6=x@vb0Pk78_0ydI9RA-Q!UH}5v*nK)Kw2Xm4gdfE delta 77 zcmV-T0J8sppaX!Q1F!~w7U2L7_YdU{$Pcy;MGlw_b`C2I_6^~)5fH@