⚡ Bolt: Replace (X**2).sum(axis) with np.einsum for faster squared norms

snapvec/__init__.py

@@ -21,11 +21,11 @@
 
 __version__ = "0.11.1"
 __all__ = [
-    "SnapIndex",
-    "PQSnapIndex",
     "IVFPQSnapIndex",
+    "PQSnapIndex",
     "ResidualSnapIndex",
+    "SnapIndex",
     "get_codebook",
-    "rht",
     "padded_dim",
+    "rht",
 ]

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],

snapvec/_file_format.py

@@ -29,10 +29,12 @@
 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
 
+from typing_extensions import Self
 
 _TRAILER_MAGIC = b"CRC2"
 _TRAILER_SIZE = 8  # 4 bytes magic + 4 bytes uint32 CRC
@@ -79,7 +81,7 @@ def finalise(self) -> None:
         self._f.write(struct.pack("<I", self._crc & 0xFFFFFFFF))
         self._finalised = True
 
-    def __enter__(self) -> "ChecksumWriter":
+    def __enter__(self) -> Self:
         return self
 
     def __exit__(
@@ -163,16 +165,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",
 ]

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("<IIIIII", _VERSION, self.dim, self.bits, self.seed, n, flags))
             f.write(struct.pack("<I", len(packed)))
@@ -599,7 +599,7 @@ def _write(f: "ChecksumWriter") -> 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.

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)
+        # np.einsum('ijk,ijk->ij', arr, arr) avoids intermediate array allocation
+        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)
@@ -441,8 +442,9 @@ def add_batch(
             for j in range(self.M):
                 Rj = residuals[:, j * self._d_sub : (j + 1) * self._d_sub]
                 # ‖R - c_j,k‖² = ‖R‖² − 2 R · c + ‖c‖²
+                # np.einsum('ij,ij->i', arr, arr) avoids intermediate array allocation
                 d2 = (
-                    (Rj * Rj).sum(1, keepdims=True)
+                    np.einsum('ij,ij->i', Rj, Rj)[:, None]
                     - 2 * Rj @ cb_T[j]
                     + cb_norms[j][None, :]
                 )
@@ -996,7 +998,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,)
+        # np.einsum('ij,ij->i', arr, arr) avoids intermediate array allocation
+        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(
@@ -1127,7 +1130,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(
@@ -1170,7 +1173,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:

snapvec/_kmeans.py

@@ -28,13 +28,16 @@ def kmeans_pp_init(
     """
     n = X.shape[0]
     centers = [X[int(rng.integers(n))]]
-    d2 = ((X - centers[0]) ** 2).sum(1)
+    diff0 = X - centers[0]
+    # np.einsum('ij,ij->i', arr, arr) avoids intermediate array allocation
+    d2 = np.einsum('ij,ij->i', diff0, diff0)
     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))
+        diff_last = X - centers[-1]
+        d2 = np.minimum(d2, np.einsum('ij,ij->i', diff_last, diff_last))
     return np.stack(centers).astype(np.float32)
 
 
@@ -50,9 +53,10 @@ def kmeans_mse(
     """
     rng = np.random.default_rng(seed)
     C = kmeans_pp_init(X, K, rng)
-    x_sq = (X ** 2).sum(1, keepdims=True)
+    # np.einsum('ij,ij->i', arr, arr) avoids intermediate array allocation
+    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, :]
+        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 +92,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, :]
+    # np.einsum('ij,ij->i', arr, arr) avoids intermediate array allocation
+    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 +117,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.
+    # np.einsum('ij,ij->i', arr, arr) avoids intermediate array allocation
     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),
     )
 
 
@@ -199,9 +205,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",
 ]

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]
+            # np.einsum('ij,ij->i', arr, arr) avoids intermediate array allocation
             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)
 
@@ -426,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(
@@ -459,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:

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("<IIIIIIII", _VERSION, self.dim, self.b1,
                                 self.b2, self.seed, n, flags, self._pdim))
@@ -318,7 +317,7 @@ def _write(f: "ChecksumWriter") -> 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:

tests/test_adversarial.py

@@ -11,7 +11,6 @@
 
 from snapvec import IVFPQSnapIndex, PQSnapIndex, ResidualSnapIndex, SnapIndex
 
-
 # --------------------------------------------------------------------------- #
 # Empty index                                                                  #
 # --------------------------------------------------------------------------- #

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,

tests/test_properties.py

@@ -16,7 +16,6 @@
 
 from snapvec import IVFPQSnapIndex, PQSnapIndex, SnapIndex
 
-
 PROFILE = settings(
     max_examples=25,
     deadline=None,

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)