feat(mps): drive SRC adaptive loop with incremental QR and batched sketching - #477
Merged
Conversation
einsum with three or more operands returned its result in the internal first-appearance order of the surviving indices instead of the order the notation requested: the pairwise recursion used that internal labeling for the last step too, and no final permutation followed. Whenever the two orders differed the returned tensor carried the wrong shape and data, silently breaking downstream axis bookkeeping. Derive the last step's output notation from the caller's requested order instead. The index sets coincide either way -- every index surviving the last step is a final-output index and vice versa -- and einsum_pair already reorders its GEMM output to an arbitrary requested order, so the reorder fuses into the step that was already running. Two-operand calls were never affected, and the existing N-operand tests only requested orders that happened to coincide with the first-appearance order, which is why the gap stayed invisible. Closes #474
IncrementalQr grows a thin QR factorization one column block at a time without refactorizing the columns already absorbed, and optionally maintains the inverse triangular factor's squared row norms alongside it -- the quantity randomized leave-one-out error estimators consume. Each append costs a block update instead of the full recompute plus inversion a from-scratch scheme pays every round. Each append runs block classical Gram-Schmidt with one reorthogonalization pass; the inverse rides the block-triangular inverse identity. Rank deficiency is detected from the maintained diagonal and terminates the factorization before the singular factor can be inverted. The accumulated basis is not unconditionally orthonormal: beyond the overt rank-deficient case, orthogonality can degrade gradually across appends whose projected parts are ill-conditioned, a regime the diagonal test cannot see because the offending entries sit legitimately above its tolerance. into_orthonormal_q therefore owns the repair, running one plain QR whenever more than a single block was appended, so callers cannot inherit the caveat by omission. Mid-layer: dense-only, consumed by ariadnetor-mps, not re-exported through the umbrella.
…etching Each growth round of the successive-randomized-compression sweep restacked every panel, refactorized the whole sketch, and inverted the resulting triangular factor -- work that grows with everything sketched so far rather than with the columns just added. Sketch columns were also recursed through the chain one at a time, so each round paid a separate tensordot chain per column per site. Feed the per-site panels to an IncrementalQr instead, so a round costs a column-block append and an incremental error-estimate update, and create and recurse columns in blocks, so each site's environments come out of one batched contraction. Environments live in per-site buffers that grow geometrically and are released once the sweep passes their last reader, which also bounds peak memory by the live prefix instead of the whole chain. Measured on a Heisenberg MPO times a random MPS (12 sites, chi = 16): the adaptive sweep runs ~20% faster, and the fixed-rank sweep ~15% faster from the batched recursion alone. Small problems (8 sites, chi = 8) are ~15% slower -- the per-round bookkeeping is not amortized there -- so the win arrives as the sketch grows. The Gaussians are now drawn per site for a whole block rather than per column across sites, so a given seed produces a different sketch than before. The sketch's distribution is unchanged and run-to-run determinism for equal seeds -- all the seed field promises -- still holds. Closes #473
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (10)
📝 WalkthroughWalkthroughSuccessiveRandomized gains batched environment sketching and incremental QR updates, with rank-deficiency handling and new regression coverage. Multi-operand einsum output ordering is corrected and tested. A Criterion benchmark compares adaptive and fixed application modes. ChangesIncremental QR foundation
SuccessiveRandomized integration
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related issues
Poem
🚥 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 |
There was a problem hiding this comment.
Pull request overview
This PR improves performance and correctness of the successive randomized compression (SRC) MPO–MPS apply path by (1) driving adaptive sketch growth with an incremental QR that updates state per appended column block, (2) batching sketch-column environment recursions per site, and (3) fixing einsum_with_backend for 3+ operands to honor the requested output index order.
Changes:
- Added
ariadnetor-linalg::IncrementalQr(+ tests) to support block-appended thin QR with optionalR^{-1}row-norm tracking for SRC’s leave-one-out estimator. - Reworked SRC internals to generate/recurse sketch columns in blocks and feed per-site sketch panels into
IncrementalQr, reducing per-round recomputation. - Fixed multi-operand einsum output ordering and added regression tests + a Criterion bench covering multi-append growth behavior.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| crates/ariadnetor-mps/tests/mps/apply_successive_randomized.rs | Adds SRC regression tests for scale invariance and late rank-deficiency re-orthonormalization. |
| crates/ariadnetor-mps/src/types.rs | Updates ApplyMethod::SuccessiveRandomized docs to reflect batched environment recursions. |
| crates/ariadnetor-mps/src/apply/successive_randomized.rs | Rewrites SRC sweep internals to use batched sketching + IncrementalQr-driven adaptive loop. |
| crates/ariadnetor-linalg/tests/einsum.rs | Adds regression test ensuring multi-operand einsum honors requested output order. |
| crates/ariadnetor-linalg/src/lib.rs | Wires new incremental_qr module and re-exports IncrementalQr/QrAppendOutcome. |
| crates/ariadnetor-linalg/src/incremental_qr/tests.rs | New unit test suite for incremental QR behavior, rank-deficiency, and error paths. |
| crates/ariadnetor-linalg/src/incremental_qr.rs | New incremental thin-QR implementation with optional inverse tracking and rank test. |
| crates/ariadnetor-linalg/src/einsum.rs | Fixes multi-operand einsum recursion to emit final output in caller-requested index order. |
| crates/ariadnetor-algorithms/Cargo.toml | Registers new mpo_mps_apply benchmark target. |
| crates/ariadnetor-algorithms/benches/mpo_mps_apply.rs | Adds end-to-end Criterion benchmark for SRC adaptive vs fixed rank. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
This was referenced Jul 15, 2026
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
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.
Summary
The successive-randomized-compression sweep compresses each bond by factorizing a sketch panel — the matrix that the sketch columns sample the partially compressed product into at that cut — and growing the sketch until an error estimate is satisfied. Every growth round restacked all panels sketched so far, refactorized them, and inverted the resulting triangular factor, so a round cost as much as the whole sketch rather than as much as the columns it added. Columns were also recursed through the chain one at a time, one tensordot chain per column per site.
This drives the per-site loop with a new incremental QR instead — a round appends a column block and updates the error estimate in place — and creates and recurses columns in blocks, so each site's environments come out of one batched contraction.
Closes #473 Closes #474
Changes
crates/ariadnetor-linalg/src/incremental_qr.rs(new):IncrementalQrgrows a thin QR one column block at a time via block Gram-Schmidt with a reorthogonalization pass, taking the backend per method like the rest of the crate's dense entry points. It optionally maintains the squared row norms ofR^-1— the quantity SRC's leave-one-out stopping rule reads — through the identity[[R, C], [0, R22]]^-1 = [[G, -G C G22], [0, G22]], so a round updates them instead of inverting from scratch. Rank deficiency is detected from the maintained diagonal and terminates the factorization before the singular factor is inverted.into_orthonormal_qre-runs one plain QR whenever more than a single block was appended.crates/ariadnetor-linalg/src/einsum.rs: fix einsum with three or more operands returning its result in an internal index order rather than the requested one. The pairwise recursion used its own first-appearance labeling for the last step too and no final permutation followed, so whenever the two orders differed the returned tensor carried the wrong shape and data. Two-operand calls were unaffected, and every pre-existing multi-operand call site and test requests an order that coincides with the first-appearance one (rgover the workspace), which is why the gap stayed invisible until the new sketch contractions needed a differing order.crates/ariadnetor-mps/src/apply/successive_randomized.rs: feed the per-site panels to anIncrementalQr; draw and recurse Gaussians per site for a block of columns; hold environments in per-site buffers that grow geometrically and are released once the sweep passes their last reader, bounding peak memory by the live prefix rather than the whole chain.crates/ariadnetor-algorithms/benches/mpo_mps_apply.rs(new): criterion bench over a Heisenberg MPO times a random MPS, with adaptive and fixed-rank arms. Its setup runs one untimed apply and asserts the adaptive arm grows past the initial sketch, so the timed loop provably exercises the multi-append path.Impact
apply_with_method, andApplyMethod/SuccessiveRandomizedParamsare untouched.IncrementalQrandQrAppendOutcomeare newpubitems inariadnetor-linalg, consumed directly byariadnetor-mpsand not re-exported through theariadnetorumbrella — the Mid-layer placement of CONTRIBUTING.md's public-API taxonomy.Test plan
Each of the four behaviors below was checked to fail on the corresponding broken implementation, by hand-reverting the fix and re-running:
f64while the norm itself stays representable; the selected bonds must match the unit-scale run's. Accumulating squared norms instead of chaininghypotcollapses the first bond from 3 to 1 and fails it.Other coverage:
IncrementalQrunit tests (f64 and Complex64): multi-append equivalence against one full QR of the stacked blocks (diagonal magnitudes, orthonormality, span, and inverse row norms against a direct inversion), first-append reduction to a plain QR, rank deficiency in a first block, the guard rejecting a backend whose memory order differs from the factorization's, post-termination panic, inverse tracking off, and every input-validation path.cargo make gate(format, clippy over all targets with warnings denied, and unit / integration / doc tests) passes.Notes
cargo bench -p ariadnetor-algorithms --bench mpo_mps_apply, 30 samples per arm: 12 sites at bond dimension 16 improves ~20% adaptive and ~15% fixed-rank, while 8 sites at bond dimension 8 regresses ~15% adaptive — the per-round bookkeeping is not amortized at that size, so the win arrives as the sketch grows. The baseline is that same bench file copied onto amainworktree, since it only calls API that already exists there. Eliminate the per-round copies in the SRC sweep #475 tracks eliminating the per-round copies of the basis and the maintained inverse, which should move the crossover down.R^-1, which overflow for states scaled below roughly 1e-154, collapsing the estimate to zero so the sweep stops at the initial sketch. The limitation predates this PR — the estimator it replaces squared the same quantities — and Leave-one-out estimate collapses to zero for extreme state scales #476 tracks the fix.seedfield promises — still holds.Summary by CodeRabbit
New Features
Bug Fixes
Benchmarks & Tests
@coderabbitai ignore
Plan-vs-actual delta
Against #473's Scope and Acceptance:
einsum_with_backendas it stood. Its first batched call returned a transposed tensor — the ordering bug above — so this PR fixes the entry point rather than routing the sketch around it, since the sketch's contractions are its first callers to need a differing order.RankDeficientappends only. That is insufficient: orthogonality also degrades gradually across full-rank appends whose projected parts are ill-conditioned, which ordinary Gaussian panels reach once a tightcutoffor a per-site cap pushes the sketch into the product's numerical noise floor, and which the diagonal test cannot see. It now runs for every multi-append history, at one QR per site.norm()call had, so it chainshypotinstead.IncrementalQr's stored basis and inverse; they are rebuilt per append instead, because the projection contraction materializes the whole basis every round regardless, so the buffer bought no asymptotic change while adding a fill invariant to maintain. The copies it would have saved are Eliminate the per-round copies in the SRC sweep #475. The environment buffers do use the specified doubling.