⚡ Bolt: Use einsum for squared norms - #185
Conversation
Replaced instances of (X ** 2).sum(axis) with np.einsum('ij,ij->i', X, X) across the codebase (specifically in kmeans, ivfpq, and pq submodules) to avoid allocating large intermediate arrays and improve overall computational throughput by approximately 3-4x.
Co-authored-by: stffns <70039235+stffns@users.noreply.github.com>
|
👋 Jules, reporting for duty! I'm here to lend a hand with this pull request. When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down. I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job! For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with New to Jules? Learn more at jules.google/docs. For security, I will only act on instructions from the user who triggered this task. |
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
|
Warning Review limit reached
Next review available in: 35 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (10)
📝 WalkthroughWalkthroughChangesThe PR replaces repeated elementwise squared-norm reductions with NumPy norm computation optimization
Estimated code review effort: 2 (Simple) | ~10 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Warning Review limit reached
Next review available in: 35 minutes You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (14)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbiteu[bot]](https://avatars.githubusercontent.com/in/3307191?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
test_perf_3.py (1)
11-26: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise the iterative
kmeans_pp_initupdate.This benchmark computes only the first-center distance. It does not execute the changed
np.minimum(d2, ...)update for subsequent centers. Add multiple center updates and compare the final distance array and selected indices between implementations.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test_perf_3.py` around lines 11 - 26, Extend the benchmark loops in test_perf_3.py beyond the initial center so they exercise iterative kmeans_pp_init distance updates using the existing d2 and d2_e arrays. Apply equivalent np.minimum updates for multiple selected centers in both implementations, then compare the final distance arrays and the selected center indices, not just the initial distances.test_perf.py (1)
1-24: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueGuard benchmark execution from imports and explicit collection.
The repository’s default pytest configuration searches only
tests, so it does not collect these root-level files by default. The files still execute all benchmark loops when imported or explicitly collected. Add amain()guard or rename the files to a benchmark-only pattern.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test_perf.py` around lines 1 - 24, Wrap the benchmark setup, timing loops, and output in a main() function, then invoke it only under an if __name__ == "__main__" guard in test_perf.py. Keep imports and constant definitions safe for collection without executing benchmark work.snapvec/_kmeans.py (1)
31-41: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd assignment-level equivalence tests.
ndarray.sumandnp.einsumcan produce different float32 reductions. Near ties or very small values, this can changeargmin, k-means++ center selection, probe ranking, or PQ codes. Compare the optimized paths with the original formulas using near-tie, tiny-value, and non-contiguous float32 inputs. Coversnapvec/_kmeans.py,snapvec/_pq.py, andsnapvec/_ivfpq.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@snapvec/_kmeans.py` around lines 31 - 41, Add assignment-level equivalence tests for the optimized distance and reduction paths in _kmeans.py, _pq.py, and _ivfpq.py, comparing them against the original ndarray.sum-based formulas. Include non-contiguous float32 inputs plus near-tie and tiny-value cases, and assert identical argmin, center-selection, probe-ranking, and PQ-code assignments where applicable.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@test_perf.py`:
- Line 24: Replace the boolean-only print validation around cb_norms1 and
cb_norms2 with an assertion or equivalent non-zero exit when np.allclose returns
false, while preserving successful completion for matching results. Apply the
same validation behavior in test_perf_2.py through test_perf_7.py.
---
Nitpick comments:
In `@snapvec/_kmeans.py`:
- Around line 31-41: Add assignment-level equivalence tests for the optimized
distance and reduction paths in _kmeans.py, _pq.py, and _ivfpq.py, comparing
them against the original ndarray.sum-based formulas. Include non-contiguous
float32 inputs plus near-tie and tiny-value cases, and assert identical argmin,
center-selection, probe-ranking, and PQ-code assignments where applicable.
In `@test_perf_3.py`:
- Around line 11-26: Extend the benchmark loops in test_perf_3.py beyond the
initial center so they exercise iterative kmeans_pp_init distance updates using
the existing d2 and d2_e arrays. Apply equivalent np.minimum updates for
multiple selected centers in both implementations, then compare the final
distance arrays and the selected center indices, not just the initial distances.
In `@test_perf.py`:
- Around line 1-24: Wrap the benchmark setup, timing loops, and output in a
main() function, then invoke it only under an if __name__ == "__main__" guard in
test_perf.py. Keep imports and constant definitions safe for collection without
executing benchmark work.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 812edc64-ca81-4516-8672-a2bf737a59d2
📒 Files selected for processing (11)
.jules/bolt.mdsnapvec/_ivfpq.pysnapvec/_kmeans.pysnapvec/_pq.pytest_perf.pytest_perf_2.pytest_perf_3.pytest_perf_4.pytest_perf_5.pytest_perf_6.pytest_perf_7.py
| t1 = time.time() | ||
| print(f"Time einsum: {t1 - t0}") | ||
|
|
||
| print(np.allclose(cb_norms1, cb_norms2)) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Fail the validation when results differ.
print(np.allclose(cb_norms1, cb_norms2)) reports False but still exits successfully. If this script is used for equivalence validation, a regression can pass automation. Raise an AssertionError or exit with a non-zero status. Apply the same change to test_perf_2.py through test_perf_7.py.
Proposed fix
-print(np.allclose(cb_norms1, cb_norms2))
+if not np.allclose(cb_norms1, cb_norms2):
+ raise AssertionError("Norm implementations differ")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| print(np.allclose(cb_norms1, cb_norms2)) | |
| if not np.allclose(cb_norms1, cb_norms2): | |
| raise AssertionError("Norm implementations differ") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@test_perf.py` at line 24, Replace the boolean-only print validation around
cb_norms1 and cb_norms2 with an assertion or equivalent non-zero exit when
np.allclose returns false, while preserving successful completion for matching
results. Apply the same validation behavior in test_perf_2.py through
test_perf_7.py.
Applied ruff automatic fixes and resolved lint/type issues discovered in CI: - Removed quotes from type annotations for self-referential types and updated to use Python 3.10+ class names as references or `__future__` imports - Formatted `__all__` arrays - Sorted imports via isort - Replaced redundant `dict()` usage in parametrization arrays in tests - Correctly combined `with` statements where possible. Co-authored-by: stffns <70039235+stffns@users.noreply.github.com>
Recent numpy releases (>=2.5.0) use the `type` statement in `__init__.pyi`, which mypy currently fails to parse when its configured `python_version` is below 3.12 (it is 3.10 in this repo). Pinned numpy in the GitHub Actions `lint` job so it successfully downloads an older numpy before running `mypy --strict`. Co-authored-by: stffns <70039235+stffns@users.noreply.github.com>
1 participant
💡 What
Replaced row-wise squared Euclidean norm calculations (
(X ** 2).sum(axis)) withnp.einsumequivalents acrosssnapvec/_kmeans.py,snapvec/_ivfpq.py, andsnapvec/_pq.py. Emulatedkeepdims=Truebehavior by appending[:, None].🎯 Why
Calculating
(X ** 2)or(X * X)creates a large intermediate NumPy array equal to the size ofXbefore performing the sum. For large N and D configurations (e.g., during IVFPQ codebook routing, K-means updates), this consumes significant memory bandwidth and incurs unnecessary allocations.📊 Impact
Prevents large temporary array allocations during critical scoring, clustering, and routing operations. Benchmarks indicate a ~3x performance improvement in the respective computational blocks compared to the original
.sum()approach.🔬 Measurement
Review tests passing and verify performance via isolated python benchmarking scripts matching the logic implemented. Expected result is ~3x speedup on operations such as
kmeans_pp_init,assign_l2, IVFPQ code additions, and PQ updates.PR created automatically by Jules for task 18326153558401829078 started by @stffns
Summary by CodeRabbit
Performance Improvements
Bug Fixes
Tests