python-math #48: feat(math): vectorize legacy kmeans hot paths (item 9a) - #2679
Draft
jucor wants to merge 1 commit into
Draft
python-math #48: feat(math): vectorize legacy kmeans hot paths (item 9a)#2679jucor wants to merge 1 commit into
jucor wants to merge 1 commit into
Conversation
This was referenced Jul 27, 2026
Draft
Draft
Draft
Draft
Draft
Draft
Draft
This was referenced Jul 27, 2026
Draft
Draft
python-math #43: docs(delphi): s7 goal state + journal — mode collapse executed, battery 20/20
#2672
Draft
Draft
Draft
Draft
This was referenced Jul 28, 2026
Draft
jucor
changed the title
3 tasks
## What
Replace the per-pair Python `_euclidean` scans in `legacy_kmeans.py` (the quirk-for-quirk port of the Clojure engine's k-means) with per-center BLAS distance columns — BIT-IDENTICAL outputs, certified twice by the full certification battery (which replays 20 recorded dataset entries through the engine and compares against Clojure reference recordings). The bit-identity plan ("Plan A") held: no tie-break divergence appeared, so no fallback was needed.
## Why
Hot paths (cProfile at 4000 participants x 300 comments x 240k votes, 112.5s total): `_euclidean` 26.4M calls / 61.8s cumulative; `np.array_equal` 8.0M calls / 38.9s cumulative via the O(n^2) `n_distinct_rows` scan; `cluster_step` 65.3s cumulative.
## Bit-identity engineering
Empirically probed on numpy 1.26.4 / OpenBLAS 0.3.23 arm64:
- REJECTED: `X @ c` (dgemv) — it reassociates relative to `float(np.dot(row, c))` for d>=4; `einsum` and `(m*m).sum(1)` differ even at d=2. Last-ulp differences are amplified by quirk Q11 — the legacy distance formula d^2 = |a|^2+|b|^2-2ab cancels catastrophically, so near-coincident points tie at EXACTLY 0.0 — into changed 0.0-ties, i.e. changed cluster lineage.
- SHIPPED: batched matmul — `(n,1,d)@(n,d,1)` for row norms and `(n,1,d)@(d,1)` for cross products — bit-equal to `float(np.dot(...))` on all 85 probed shape/scale combos (n=1..33422, d=1..783, scales 1e-8/1/1e8, C and F memory order), including the vw dataset's knife-edge pair (both distances EXACTLY 0.0) and NaN propagation.
- The column combine keeps the scalar code's exact order and associativity: `(row_norms + |c|^2) - 2.0*cross`; the floor is `np.where(d2 < 0.0, 0.0, d2)` so NaN propagates (never a maximum-clamp); same IEEE sqrt.
## Changes (`legacy_kmeans.py` only)
- New `_row_norms` + `_euclidean_col` (bit-identical distance columns).
- `cluster_step`: the columns are folded in the existing Clojure scan order (hash order for >8 clusters, input order for <=8) with the scalar rule `d <= best` (ties go to the LATER cluster; NaN never wins); members are regrouped by ascending row index, which equals the scalar append order.
- `most_distal`: same inner fold; exact outer last-wins reduction (a NaN first row sticks; later NaN rows are skipped; otherwise last argmax).
- `n_distinct_rows(bound=)`: vectorized elimination passes with `array_equal(..., equal_nan=True)` semantics; `clean_start_clusters` passes `bound=k` so `min(k, .)` stays exact without the O(n^2) count.
- `_recenter_center`: index-gather (same values, same mean).
- `_euclidean` / `same_clustering` / `weighted_mean` semantics unchanged.
## Testing (TDD)
RED = 4 ImportError pins (`_euclidean_col` / `_row_norms`) plus a TypeError pin for the `bound` kwarg. 11 new tests: bit-equality against a VERBATIM `_scalar_d2_reference` copy of the pre-vectorization formula (exact ==), the knife-edge 0.0 case, NaN propagation, `cluster_step` / `most_distal` equivalence against verbatim scalar reference loops (grid-tie fixtures, the >8 and <=8 scan orders, weights, k>n, NaN), and bounded-distinct semantics (NaN, -0.0 == 0.0).
## Results
- Full suite: 1171 passed / 22 skipped / 44 xfailed / 2 xpassed (baseline 1160 + 11 new; zero new failures). Pyright clean.
- Battery: certify 20 entries — MATCH=20, DIVERGENCE=0, SKIPPED=0, ERROR=0 — TWICE (on the vectorized tree; final bytes after an annotation-only cleanup).
- Bench at 8000 x 400 x 480k votes: cold 68.53s → 3.58s (19x), warm 159.53s → 3.77s (42x).
- Bench at the FULL 33,422 x 783 / 2.0M-vote shape: cold 28.15s, warm 26.66s — versus ~31 min warm before (~70x); the largest production conversation now ticks in under 30 seconds.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
commit-id:52a753ee
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Replace the per-pair Python
_euclideanscans inlegacy_kmeans.py(the quirk-for-quirk port of the Clojure engine's k-means) with per-center BLAS distance columns — BIT-IDENTICAL outputs, certified twice by the full certification battery (which replays 20 recorded dataset entries through the engine and compares against Clojure reference recordings). The bit-identity plan ("Plan A") held: no tie-break divergence appeared, so no fallback was needed.Why
Hot paths (cProfile at 4000 participants x 300 comments x 240k votes, 112.5s total):
_euclidean26.4M calls / 61.8s cumulative;np.array_equal8.0M calls / 38.9s cumulative via the O(n^2)n_distinct_rowsscan;cluster_step65.3s cumulative.Bit-identity engineering
Empirically probed on numpy 1.26.4 / OpenBLAS 0.3.23 arm64:
X @ c(dgemv) — it reassociates relative tofloat(np.dot(row, c))for d>=4;einsumand(m*m).sum(1)differ even at d=2. Last-ulp differences are amplified by quirk Q11 — the legacy distance formula d^2 = |a|^2+|b|^2-2ab cancels catastrophically, so near-coincident points tie at EXACTLY 0.0 — into changed 0.0-ties, i.e. changed cluster lineage.(n,1,d)@(n,d,1)for row norms and(n,1,d)@(d,1)for cross products — bit-equal tofloat(np.dot(...))on all 85 probed shape/scale combos (n=1..33422, d=1..783, scales 1e-8/1/1e8, C and F memory order), including the vw dataset's knife-edge pair (both distances EXACTLY 0.0) and NaN propagation.(row_norms + |c|^2) - 2.0*cross; the floor isnp.where(d2 < 0.0, 0.0, d2)so NaN propagates (never a maximum-clamp); same IEEE sqrt.Changes (
legacy_kmeans.pyonly)_row_norms+_euclidean_col(bit-identical distance columns).cluster_step: the columns are folded in the existing Clojure scan order (hash order for >8 clusters, input order for <=8) with the scalar ruled <= best(ties go to the LATER cluster; NaN never wins); members are regrouped by ascending row index, which equals the scalar append order.most_distal: same inner fold; exact outer last-wins reduction (a NaN first row sticks; later NaN rows are skipped; otherwise last argmax).n_distinct_rows(bound=): vectorized elimination passes witharray_equal(..., equal_nan=True)semantics;clean_start_clusterspassesbound=ksomin(k, .)stays exact without the O(n^2) count._recenter_center: index-gather (same values, same mean)._euclidean/same_clustering/weighted_meansemantics unchanged.Testing (TDD)
RED = 4 ImportError pins (
_euclidean_col/_row_norms) plus a TypeError pin for theboundkwarg. 11 new tests: bit-equality against a VERBATIM_scalar_d2_referencecopy of the pre-vectorization formula (exact ==), the knife-edge 0.0 case, NaN propagation,cluster_step/most_distalequivalence against verbatim scalar reference loops (grid-tie fixtures, the >8 and <=8 scan orders, weights, k>n, NaN), and bounded-distinct semantics (NaN, -0.0 == 0.0).Results
Co-Authored-By: Claude Fable 5 noreply@anthropic.com
commit-id:52a753ee
Stack: