Lift the >2-tensor limit on external auxiliary hyperindices (TNv3 canonicalization + BTAS batched eval)#574
Merged
Conversation
TensorNetworkV3::canonicalize_slots classified each named (external) index by its edge's vertex count, asserting the count was 1 (an open bra/ket/aux slot) or 2 (an internal contraction edge that is also named). A named index shared among more than two tensor slots -- a high-order hyperindex, e.g. an auxiliary/batching index riding in the aux slot of every factor as in Laplace-transform MP2 or tensor hypercontraction -- forms an edge with >2 vertices and tripped `SEQUANT_ASSERT(edge_it->vertex_count() == 2)`. This made `binarize<EvalExprBTAS>(deserialize<ResultExpr>(...))` throw for any expression whose result carries such an index (whereas hyperindices that are fully contracted are anonymous and already handled, since the named-index loop skips them). Classify a >2-vertex named edge whose vertices are all aux as a TensorAux slot; a non-aux high-order hyperindex remains unsupported and now fails with an explanatory assertion message rather than a bare count check. Adds a regression test reproducing the Laplace-MP2 energy expression that verifies the eval tree builds and carries the batching index z1 in the aux slot of every tensor node. NB this covers tree construction only; numerically evaluating such a tree is a separate concern -- every node is a batched (Hadamard) contraction over z1, which the BTAS backend's btas::contract does not support (an index in both operands is always summed, never batched).
btas::contract classifies every index as free-left, free-right or contracted --
an index shared by both operands is always summed -- so it cannot evaluate a
contraction that carries a batch axis: a label present in both operands AND the
result, e.g. an auxiliary/quadrature index ridden through the contraction as in
Laplace-transform MP2 or tensor hypercontraction. Such a product previously
aborted inside btas::contract.
ResultTensorBTAS::prod now detects batch axes (batch_indices: labels common to
left, right and result) and, when present, routes through batched_contract():
the batch axes are permuted to the front of both operands (making each batch
slice a contiguous block in row-major storage) and each slice is contracted
with plain btas::contract, or -- for the degenerate slice shapes -- a dot
(empty result remainder) or a scaling (a scalar operand slice). The per-slice
results are assembled and permuted back into the requested result order.
Tests (test_eval_btas.cpp):
- btas_batched_contract: unit-tests prod() directly on every slice shape
(contract+batch, dot+batch, scalar*tensor+batch, non-canonical result order).
- eval_btas_batched_over_aux: end-to-end binarize + evaluate of
R{;;z1} = A{;a1;z1} B{a1;a2;z1} C{a2;;z1} (z1 shared by 3 tensors and carried
into the result) against a hand-computed reference.
…or clean when assertions are disabled (SEQUANT_ASSERT_BEHAVIOR=IGNORE)
…-end
Adds a case to the "non-covariant indices" section evaluating
R{;;x1} = A{;a1;x1} B{a1;a2;x1} C{a2;;x1}, where the auxiliary index x1 is
shared by 3 tensors AND carried into the result -- the >2-tensor named
hyperedge unlocked by the TNv3 canonicalize_slots fix. TiledArray evaluates the
resulting batched-over-x1 contraction natively (einsum keeps an index common to
both operands and the result), matching a hand-written einsum reference. No
backend change is needed for TA (unlike BTAS, whose btas::contract required the
batched_contract path).
The enumerator had no users; a high-order auxiliary hyperindex is classified as TensorAux (see canonicalize_slots), not a distinct aux-aux bundle type.
Contributor
There was a problem hiding this comment.
Pull request overview
Adds support for external high-order (>2-tensor) auxiliary hyperindices by updating TNv3 slot canonicalization and extending the BTAS backend to evaluate batched (Hadamard) contractions over shared aux axes.
Changes:
- TNv3
canonicalize_slotsnow recognizes named hyperedges with >2 vertices asTensorAuxwhen all incident slots are auxiliary. - BTAS backend detects batch axes (labels present in left, right, and result) and evaluates them via per-batch slicing + contraction.
- Adds TA + BTAS unit/e2e tests covering high-order aux hyperindices and BTAS batched contraction behavior.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_eval_ta.cpp | Adds TA regression test for a >2-tensor external aux hyperindex carried into the result. |
| tests/unit/test_eval_btas.cpp | Adds BTAS tests for TNv3 tree building and BTAS batched-over-aux evaluation paths. |
| SeQuant/core/tensor_network/v3.cpp | Extends slot classification to handle named >2-vertex aux hyperedges. |
| SeQuant/core/tensor_network/slot.hpp | Removes TensorAuxAux from IndexSlotType. |
| SeQuant/core/eval/backends/btas/result.hpp | Implements batch-axis detection and batched_contract() for BTAS tensor products. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+805
to
+815
| [[maybe_unused]] bool all_aux = true; | ||
| for (std::size_t v = 0; v != edge_it->vertex_count(); ++v) | ||
| if (edge_it->vertex(v).getOrigin() != Origin::Aux) { | ||
| all_aux = false; | ||
| break; | ||
| } | ||
| SEQUANT_ASSERT(all_aux && | ||
| "TensorNetworkV3::canonicalize_slots: high-order " | ||
| "(shared among >2 tensor slots) non-auxiliary " | ||
| "hyperindices are not supported"); | ||
| slot_type = IndexSlotType::TensorAux; |
Comment on lines
+189
to
+202
| auto const nb = batch.size(); | ||
| auto tail_extents = [](T const& t, std::size_t skip) { | ||
| container::svector<std::size_t> e; | ||
| std::size_t sz = 1; | ||
| for (std::size_t i = skip; i != t.rank(); ++i) { | ||
| e.push_back(t.extent(i)); | ||
| sz *= t.extent(i); | ||
| } | ||
| return std::pair{std::move(e), sz}; | ||
| }; | ||
| std::size_t P = 1; | ||
| for (std::size_t i = 0; i != nb; ++i) P *= Lp.extent(i); | ||
| auto const [rlext, rlsz] = tail_extents(Lp, nb); | ||
| auto const [rrext, rrsz] = tail_extents(Rp, nb); |
Comment on lines
38
to
42
| /// connecting two tensor vector index slots | ||
| TensorBraKet, | ||
| /// in tensor aux index slot | ||
| TensorAux, | ||
| /// connecting two tensor aux index slots | ||
| TensorAuxAux, | ||
| }; |
- canonicalize_slots: a non-auxiliary high-order (>2-slot) hyperindex now throws unconditionally instead of via SEQUANT_ASSERT, so it hard-errors in every build config (previously, with assertions disabled it silently fell through to TensorAux). all_aux is now read unconditionally, so the [[maybe_unused]] is gone. - batched_contract: assert the two operands agree on every batch extent before deriving the batch count P from the left operand and slicing.
Member
Author
|
Thanks for the review. Addressed in 8474a9d:
|
evaleev
changed the title
evaleev
added a commit
that referenced
this pull request
Jul 10, 2026
Reconcile the predicted-peak / batching / dry-run work with master's merged #576 (outer-product pruning + build_context skip + fast_flops), #575/#574 (iterative evaluate()/binary_node), and TNv3 high-order aux work. Conflict resolutions: - options.hpp: CostParams and OptimizeOptions carry BOTH the branch's peak_threshold / term_batch_axes and master's prune_outer_products. - single_term.hpp: keep the perf-first objective dispatch (perf_first, DenseTimeSpace/DenseTimeSpaceBatched, peak_threshold, out_axes) and thread master's prune_outer_products into every model. - cost_model.hpp: PeakModel / PeakBatchedModel keep the branch's perf-first members (perf_first, peak_threshold, numeric_size, charge_batch_recompute) and gain master's prune_outer_products; build_context bodies auto-merged so the pruning skip + fast_flops precompute coexist with the perf-first tables. - optimize.cpp: CostParams init passes both peak_threshold and prune_outer_products. - eval.hpp: keep the branch's malloc_trim release_after_op() hook at each op. - test_optimize.cpp: keep both the branch's threshold/perf-first tests and master's pruning + fast-flops-parity tests. The branch's own (superseded) outer-product-pruning commits were dropped before this merge since #576 supersedes them; the predict-hook add/remove experiment nets to no hook, matching master.
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
SeQuant already supports hyperindices — indices shared among more than two tensors — as exercised by the existing tensor-hypercontraction / THC tests and the order-3
T{a1;i1} T{a2;i1} T{a3;i1}test. This PR does not add hyperindex support; it lifts a limitation that only shows up at higher order.In every existing test the hyperindex is either contracted (an internal/anonymous mode — skipped by the named-index handling in canonicalization) or external but shared by at most two tensor slots. Neither reaches the one code path that assumed a named index's edge connects at most two vertices, so the limit was never hit — the existing coverage tops out at order-3 contracted hyperindices.
Laplace-transform MP2 is the first case with an external hyperindex of order > 2: a batching/quadrature index
z1rides in the aux slot of every factor and of the resultE{;;z1}(order 9 here). That tripped the assertion, and — once past it — could not be evaluated. Both layers are addressed:1. Tree building (TNv3 canonicalization)
TensorNetworkV3::canonicalize_slotsclassified each named index by its edge's vertex count, handling 1 (an open bra/ket/aux slot) or 2 (an internal contraction that is also named, e.g. a protoindex) and assertingvertex_count() == 2otherwise. A named index shared by >2 tensor slots forms a hyperedge with >2 vertices and tripped that assert.Fix: classify an all-aux, >2-vertex named edge as a
TensorAuxslot. A non-auxiliary high-order hyperindex has no well-defined slot type and is rejected with an explicitthrow(in every build config, not only assertion-enabled ones).2. Numerical evaluation
Each contraction node over such an index is a batched (Hadamard) contraction: the index is present in the left operand, the right operand and the result, so it must be iterated, not summed.
btas::contracthas no batch category (an index in both operands is always summed), so it aborted.ResultTensorBTAS::prodnow detects batch axes (batch_indices) and routes through a newbatched_contract()that permutes the batch axes to the front (making each slice contiguous in row-major storage) and contracts each slice — falling back to a dot or a scaling for the degenerate slice shapes.einsumsupports batch axes natively. Only the layer-1 fix was required.Tests
binarize_highorder_aux_hyperindex(test_eval_btas.cpp) — reproduces the Laplace-MP2 energy expression; asserts the eval tree builds and carriesz1in the aux slot of every tensor node. Confirmed red→green against the fix.btas_batched_contract— unit-testsprod()on every slice shape (contract+batch, dot+batch, scalar·tensor+batch, non-canonical result order).eval_btas_batched_over_aux— end-to-endbinarize + evaluateof a 3-tensor batched-over-aux contraction against a hand-computed reference.test_eval_ta.cppnon-covariant section — end-to-end TA evaluation of the same class of expression against aTA::einsumreference.Full BTAS unit suite: 7194 assertions / 65 cases pass. TA
eval_with_tiledarraycase passes with the new test.Notes
IndexSlotType::TensorAuxAuxis removed — it had no users, and a high-order auxiliary hyperindex is classified asTensorAux, so there is no distinct aux-aux slot type.batched_contractbefore slicing.