add unified contraction algebra interface with tropical and bf16 applications - #125
add unified contraction algebra interface with tropical and bf16 applications#125DeanTMaxim wants to merge 16 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a generic contraction_algebra monkey-patch interface to TensorCircuit-NG, enabling custom contraction primitives such as max-plus (tropical) algebra, counting, and configuration recovery. It includes reference implementations, examples, and comprehensive test suites. The review feedback suggests improving code conciseness by using np.prod instead of explicit loops, utilizing be.abs() for clearer floating-point comparisons, refactoring redundant logic in the tropical context manager, and adopting a more robust declarative approach for test path configuration in conftest.py.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
fad6f8d to
2f7ecac
Compare
refraction-ray
left a comment
refraction-ray
left a comment
There was a problem hiding this comment.
I am just quickly going through and raise some more evident catches, the most important on is on the interface deisgn, you need to consider the node patch case together to design the interface in one run, instead of further changing the interface later
2f7ecac to
afde184
Compare
DeanTMaxim
changed the title
7bcb112 to
03a192f
Compare
…ications ContractionAlgebra + Representation ABCs (boundary encode/decode) activated via set_contractor(algebra=...); routed in-source by cons._algebraic_base_contraction (no monkey-patch). Two reference applications in applications/: tropical (max-plus ground-state energy, configuration recovery, degeneracy counting; arXiv:2008.06888) and complex<bfloat16> pair-algebra (4M real bf16 matmul; boundary encode/decode keeps the pair axis off tn.Node, dodging the axis==edge wall).
03a192f to
539ff5d
Compare
9968429 to
a1f8c1b
Compare
…ing concatenation
7612492 to
5de41f8
Compare
- Replace prefer_einsum: bool on ContractionAlgebra ABC with get_contractor_kwargs()
method (default {}), keeping the ABC clean of cotengra-specific flags
- ComplexPairAlgebra overrides to return {'prefer_einsum': True}
- cons.py unpacks via **alg.get_contractor_kwargs()
- De-hardcode float32/complex64 in _pair_to_complex to use cons.rdtypestr/dtypestr
- Remove personal pythonpath config from pyproject.toml and example docstring
- Add unit test for get_contractor_kwargs default behavior
- Tighten test_tropical_public_api_surface to direct import assertions
- Merge base.py into contraction_algebra.py (122 lines total → single module) - Remove subpackage directory (__init__.py + base.py → contraction_algebra.py) - Fix deep imports in cons.py and test_contraction_algebra.py
- Tighten GHZ-8 test tolerance (5e-2 → 1.5e-2)
- Replace fake_node string with real tn.Node in merge_single_gates test
- Fix 'Opt-in' comment → 'ALGEBRAIC EXECUTION PATH'
- Document thread-safety limitation of _aux_outputs_store
- Explain be/kbe naming and einsum diagonal safety
- Inline _expand_to_layout_pair (single call-site)
- Fix CHANGELOG: set_contractor → set_contraction_algebra
- Revert unrelated black version pin in requirements-dev.txt
- Add _stash_aux_outputs({}) to legacy contraction path
…ebra - _contraction_algebra defaults to None; _standard singleton for fallback - _ns_alg check uses is not None (not isinstance) - _algebraic_base_contraction uses alg or _standard - Zero overhead for default users: legacy path runs with no diff - Cleaner branch, better backward compat and reviewer friendliness
Codecov Report❌ Patch coverage is
... and 8 files with indirect coverage changes 🚀 New features to boost your workflow:
|
269c756 to
0184776
Compare
a64526f to
915cc76
Compare
| from .backends import get_backend | ||
| from .simplify import _multi_remove | ||
| from .contraction_algebra import ( | ||
| ContractionAlgebra as _ContractionAlgebra, |
There was a problem hiding this comment.
why as, and there is also lower case _contraction_algebra variable in this file, dangerously similar and confusing
| return _contraction_algebra | ||
|
|
||
|
|
||
| def set_contraction_algebra(alg: Optional[_ContractionAlgebra]) -> None: |
There was a problem hiding this comment.
no support for default arg, say set_contraction_algebra() with no arg should recover the default behavior
| """Decode the contraction output under a non-standard algebra. | ||
|
|
||
| ``rep.decode`` is responsible for stashing any aux internally.""" | ||
| if not ns: |
There was a problem hiding this comment.
no need for ns, if not ns, in the parent function there should be no call on _decode
| if alg is not None: | ||
| rep = alg.representation | ||
| else: | ||
| rep = IdentityRepresentation() # no-op; _decode skips it when ns=False |
| size_dict, | ||
| algorithm, | ||
| alg is not None, | ||
| algebra=alg, |
There was a problem hiding this comment.
merge the two variable? as algebra=None means ns=False
| ctg=ctg, | ||
| ) | ||
|
|
||
| final = _decode(alg is not None, rep, backend, final, output_set) |
There was a problem hiding this comment.
if outside the _decode call
|
|
||
| To change the contraction algebra, use ``cons.set_contraction_algebra(alg)`` | ||
| separately (the algebra is orthogonal to the contractor configuration). The | ||
| ``tropical()`` / ``bcomplex32()`` / ``counting_tropical()`` context managers |
There was a problem hiding this comment.
no, dont define these sub context managers, just use something like set_contraction_algebra(TropicalAlgebra)
| return tensor, {} | ||
|
|
||
|
|
||
| class ContractionAlgebra(ABC): |
There was a problem hiding this comment.
is there elegant way to merge Representation class and Algebra class?
There was a problem hiding this comment.
is it okay just define encode and decode in Algebra class?
| return be.einsum(eq, *operands) | ||
|
|
||
|
|
||
| class PairTensor: |
There was a problem hiding this comment.
why bother the class, just bare tuple suffices?
|
|
||
|
|
||
| @contextlib.contextmanager | ||
| def bcomplex32() -> Iterator[None]: |
915cc76 to
bc7671c
Compare
|
|
|
||
| # --- Task 12: end-to-end through real tc.Circuit + wall-avoidance canary --- | ||
|
|
||
| import tensorcircuit as tc |
There was a problem hiding this comment.
middle import is not allowed
| # --- Task 12: end-to-end through real tc.Circuit + wall-avoidance canary --- | ||
|
|
||
| import tensorcircuit as tc | ||
| from applications.bcomplex32_algebra import bcomplex32 |
There was a problem hiding this comment.
how can you import something from application instead of tensoricrcuit.application?
| from applications.bcomplex32_algebra import bcomplex32 | ||
|
|
||
|
|
||
| def test_bf16_end_to_end_matches_complex64(): |
There was a problem hiding this comment.
set jax, tf and numpy backend with fixtures?
| c.cnot(i, i + 1) | ||
| st = np.asarray(c.state()) | ||
| assert st.shape == (16,) # ran cleanly, no axis==edge crash | ||
| import tensorcircuit.cons as cons |
There was a problem hiding this comment.
no middle import, too ugly
| @@ -0,0 +1,470 @@ | |||
| import numpy as np | |||
| import tensorcircuit as tc | |||
| import tensorcircuit.cons as cons | |||
|
|
||
|
|
||
| def test_public_api_surface(): | ||
| from tensorcircuit import contraction_algebra as tca |
There was a problem hiding this comment.
no middle and repeated import
| def test_public_api_surface(): | ||
| from tensorcircuit import contraction_algebra as tca | ||
|
|
||
| # After Task 14 the package exports only the 4 base names; activation lives |
There was a problem hiding this comment.
dont leave think loudly type comments
| "IdentityRepresentation", | ||
| ]: | ||
| assert hasattr(tca, name), name | ||
| # The old monkey-patch API names are intentionally gone: |
There was a problem hiding this comment.
why you test something are changed to ensure they are not existing anymore, weird...
| StandardAlgebra, | ||
| Representation, | ||
| IdentityRepresentation, | ||
| ) |
| import numpy as np | ||
| import tensornetwork as tn | ||
| import opt_einsum | ||
| import pytest |
2 participants
Summary
A generic, in-source
ContractionAlgebrainterface: run a tensor-network contraction over an arbitrary algebra (kernels + a boundary representation) viatc.set_contractor(algebra=...), routed bycons._algebraic_base_contraction. Two reference applications inapplications/:Standard contractions are unchanged.
Motivation
tc-ng contractions normally run over the standard complex semiring. Some workloads need a different algebra — tropical (max-plus) for spin-glass ground states, or a low-precision storage (bf16 pair) for memory. Rather than special-case each in the contraction internals, this PR adds one generic entry point. The key choice is where a non-standard representation plugs in: the
_algebraic_base_contractionboundary is the single place raw tensors are materialized outside the kernels, so a representation can transform storage once (leaves in, final out) without touchingtn.Nodeor the topology. That keeps a non-standard axis (e.g. the bf16 pair axis) offtn.Node, which assumes axis==edge — avoiding that conflict by construction.Design
Two ABCs in
tensorcircuit/contraction_algebra/base.py:ContractionAlgebra(tensordot/einsum kernels + optional hooks) carries aRepresentation(boundaryencode/decodecodec), bundled so a kernel can't be mis-paired with a codec._algebraic_base_contractionrunsencode→kernels→decodeat the topology boundary (after_extract_topology, beforetn.Node); three small in-source guards incons.pyroute non-standard algebras and skippreprocessing. Counting's degeneracy rides an aux side-channel (degeneracy()).Breaking changes
None for upstream
master— master has nocontraction_algebramodule; the diff is additive except ~11 internal lines incons.py. Activate an algebra withtc.set_contractor(algebra=...), or thetropical()/counting_tropical()/bcomplex32()context managers.Testing
~65 tests: standard==native keystone; tropical energy / counting (scalar + non-scalar) / tracking vs brute force; bf16 4M kernel units + end-to-end vs
complex64+ an axis==edge wall canary. Author-run gates green (black,mypy --strict,pylint 10.00/10on the new modules).Notes for review
einsumrejects bf16, so_pair_einsumupcasts bf16→float32 per call (structure/quantization unchanged).ml_dtypesis currently undeclared (lazily imported by the bf16 path; not inpyproject.toml) — maintainer decision whether to add it (e.g. an optional extra).Out of scope (follow-up PRs)
GPU bf16 kernel + benchmark (L3);
set_dtype("bcomplex32")pipeline (L4); non-scalar trackingrecover_configuration(orthogonal, separately gated).