From 539ff5d1f6f1592d39a0323cb41fe55c48bbfd7d Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Fri, 17 Jul 2026 23:01:39 +0800 Subject: [PATCH 01/16] add unified contraction algebra interface with tropical and bf16 applications 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 pair-algebra (4M real bf16 matmul; boundary encode/decode keeps the pair axis off tn.Node, dodging the axis==edge wall). --- CHANGELOG.md | 4 + applications/bcomplex32_algebra.py | 104 +++ applications/tropical_algebra.py | 811 ++++++++++++++++++ examples/tropical_ising.py | 66 ++ pyproject.toml | 1 + requirements/requirements-dev.txt | 2 +- tensorcircuit/cons.py | 211 ++++- tensorcircuit/contraction_algebra/__init__.py | 27 + tensorcircuit/contraction_algebra/base.py | 95 ++ tests/_tropical_test_utils.py | 379 ++++++++ tests/test_bcomplex32_algebra.py | 90 ++ tests/test_contraction_algebra.py | 409 +++++++++ tests/test_tropical_algebra.py | 272 ++++++ tests/test_tropical_config.py | 246 ++++++ tests/test_tropical_counting.py | 364 ++++++++ tests/test_tropical_example.py | 39 + tests/test_tropical_ising.py | 223 +++++ 17 files changed, 3312 insertions(+), 31 deletions(-) create mode 100644 applications/bcomplex32_algebra.py create mode 100644 applications/tropical_algebra.py create mode 100644 examples/tropical_ising.py create mode 100644 tensorcircuit/contraction_algebra/__init__.py create mode 100644 tensorcircuit/contraction_algebra/base.py create mode 100644 tests/_tropical_test_utils.py create mode 100644 tests/test_bcomplex32_algebra.py create mode 100644 tests/test_contraction_algebra.py create mode 100644 tests/test_tropical_algebra.py create mode 100644 tests/test_tropical_config.py create mode 100644 tests/test_tropical_counting.py create mode 100644 tests/test_tropical_example.py create mode 100644 tests/test_tropical_ising.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 17d4f138..b5552469 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ ### Added +- add unified ContractionAlgebra interface (`set_contractor(algebra=...)` + boundary + encode/decode) with tropical (max-plus) and complex pair-algebra + reference applications. + - Add efficient `expectation_pss` method for `U1Circuit`. - Support MVP mode for timeevol methods. diff --git a/applications/bcomplex32_algebra.py b/applications/bcomplex32_algebra.py new file mode 100644 index 00000000..5d410124 --- /dev/null +++ b/applications/bcomplex32_algebra.py @@ -0,0 +1,104 @@ +"""complex pair-algebra — a reference APPLICATION of ContractionAlgebra. + +Pair repr: complex tensor = stack([re, im], axis=-1) of bf16. Contraction = 4 real +bf16 matmuls (4M). Activated via ``set_contractor(algebra=ComplexPairAlgebra())`` or +the ``bcomplex32()`` CM. encode/decode at the ``_algebraic_base_contraction`` boundary +keep the pair axis off tn.Node (dodges the axis==edge wall). +""" + +from typing import Any, Dict, Iterator, List, Tuple +import contextlib + +import tensorcircuit.cons as cons +from tensorcircuit.contraction_algebra import ContractionAlgebra, Representation + +Tensor = Any +Backend = Any + + +def _bf16_dtype() -> Any: + import ml_dtypes + + return ml_dtypes.bfloat16 + + +def _complex_to_pair(be: Backend, t: Tensor) -> Tensor: + """complex tensor -> stack([re, im], axis=-1) of bf16.""" + bf = _bf16_dtype() + re = be.cast(be.real(t), bf) + im = be.cast(be.imag(t), bf) + return be.stack([re, im], axis=-1) + + +def _pair_to_complex(be: Backend, pair: Tensor) -> Tensor: + """pair of bf16 -> complex64 tensor (recombine; no copy risk via cast).""" + re = be.cast(pair[..., 0], "float32") + im = be.cast(pair[..., 1], "float32") + return be.cast(re + 1j * im, "complex64") + + +def _pair_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + """Complex tensordot = 4 real bf16 tensordots (4M). Uses be.tensordot (never patched).""" + ar, ai = a[..., 0], a[..., 1] + br, bi = b[..., 0], b[..., 1] + cr = be.tensordot(ar, br, axes) - be.tensordot(ai, bi, axes) + ci = be.tensordot(ar, bi, axes) + be.tensordot(ai, br, axes) + return be.stack([cr, ci], axis=-1) + + +def _pair_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: + """Complex einsum = real bf16 einsums (4M for 2 operands, 2 for 1). + + Each pair component is upcast bf16 -> float32 before ``be.einsum``: numpy's + einsum does not support ``ml_dtypes.bfloat16`` (``TypeError: invalid data + type for einsum``), and float32 is the universal portable compute dtype for + the 4M decomposition. The bf16 quantization happened at encode time + (``_complex_to_pair``), so the values flowing in are already bf16-quantized; + this cast only widens the compute dtype, it does not undo quantization. + ``_pair_tensordot`` needs no such cast because ``np.tensordot`` accepts bf16. + """ + if len(operands) == 1: + a = operands[0] + ar = be.cast(a[..., 0], "float32") + ai = be.cast(a[..., 1], "float32") + return be.stack([be.einsum(eq, ar), be.einsum(eq, ai)], axis=-1) + a, b = operands + ar = be.cast(a[..., 0], "float32") + ai = be.cast(a[..., 1], "float32") + br = be.cast(b[..., 0], "float32") + bi = be.cast(b[..., 1], "float32") + cr = be.einsum(eq, ar, br) - be.einsum(eq, ai, bi) + ci = be.einsum(eq, ar, bi) + be.einsum(eq, ai, br) + return be.stack([cr, ci], axis=-1) + + +class PairBf16Representation(Representation): + name = "pair_bf16" + + def encode(self, be: Backend, tensors: List[Tensor]) -> List[Tensor]: + return [_complex_to_pair(be, t) for t in tensors] + + def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]: + return _pair_to_complex(be, tensor), {} + + +class ComplexPairAlgebra(ContractionAlgebra): + name = "bcomplex32_pair" + representation = PairBf16Representation() + prefer_einsum = True # pair operands carry a trailing storage axis + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + return _pair_tensordot(be, a, b, axes) + + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + return _pair_einsum(be, eq, *operands) + + +@contextlib.contextmanager +def bcomplex32() -> Iterator[None]: + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(ComplexPairAlgebra()) + try: + yield + finally: + cons.set_contraction_algebra(prev) diff --git a/applications/tropical_algebra.py b/applications/tropical_algebra.py new file mode 100644 index 00000000..a377351e --- /dev/null +++ b/applications/tropical_algebra.py @@ -0,0 +1,811 @@ +"""Tropical (max-plus) contraction algebra -- a reference APPLICATION of the generic +``ContractionAlgebra`` ABCs shipped in ``tensorcircuit.contraction_algebra``. + +The feature is ``tensorcircuit.contraction_algebra`` -- the ``ContractionAlgebra`` and +``Representation`` ABCs -- wired in-source into ``cons._base``, which routes any +non-standard algebra through ``cons._algebraic_base_contraction`` (no monkey-patching). +A custom algebra is activated either via ``tc.set_contractor(algebra=...)`` or by the +``tropical()`` / ``counting_tropical()`` context managers below. This file is one +complete algebra built on top of the ABCs -- an importable reference module, not a +runnable demo (for that, see ``examples/tropical_ising.py``). Kept under +``applications/`` as a reference application; promote it to its own package/location if +it needs active development, independent distribution, or outgrows ``applications/`` +conventions. + +Implements the three standard tropical-tensor-network outputs (energy / +configuration / degeneracy) from Liu, Wang, Zhang PRL 126, 090506 (2021) +(arXiv:2008.06888). Contract under a tropical algebra via:: + + from applications.tropical_algebra import tropical + with tropical(): + ... # max-plus contraction + +Sections (consolidated from the original package modules): + 1. max-plus primitives + MaxPlusAlgebra (was tensorcircuit/contraction_algebra/tropical.py) + 2. counting (energy, degeneracy) + split_energy_count (was .../counting.py) + 3. tracking + configuration recovery (was .../config.py) + 4. context managers: tropical(track=...) / counting_tropical() +""" + +import contextlib +import logging +from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple + +import numpy as np + +import tensorcircuit.cons as cons +from tensorcircuit.contraction_algebra import ContractionAlgebra, Representation + +Tensor = Any +Backend = Any +_EPS = 1e-9 + +logger = logging.getLogger(__name__) + + +# ===== Section 1: max-plus primitives + MaxPlusAlgebra ===== + + +def _pair_layout( + be: Backend, a: Tensor, b: Tensor, axes: Any +) -> "tuple[Tensor, Tensor, tuple[int, ...]]": + """Transpose+reshape ``a``, ``b`` into the broadcast pair layout. + + Returns ``(a3, b3, out_shape)`` where ``a3`` is shape ``(m, k, 1)``, ``b3`` is + ``(1, k, n)``, and ``out_shape = a_free_shape + b_free_shape``. Here ``m`` is the + product of ``a``'s free axes, ``n`` of ``b``'s free axes, and ``k`` of the + contracted axes. The native broadcast of ``a3 + b3`` yields the ``(m, k, n)`` + pair-sum tensor required by max-plus (and counting) pairwise contraction. + """ + ashape = tuple(int(x) for x in be.shape_tuple(a)) + bshape = tuple(int(x) for x in be.shape_tuple(b)) + if isinstance(axes, int): + a_axes = list(range(len(ashape) - axes, len(ashape))) + b_axes = list(range(0, axes)) + else: + a_axes, b_axes = list(axes[0]), list(axes[1]) + a_free = [i for i in range(len(ashape)) if i not in a_axes] + b_free = [i for i in range(len(bshape)) if i not in b_axes] + a_t = be.transpose(a, tuple(a_free + a_axes)) # contracted axes to the tail + b_t = be.transpose(b, tuple(b_axes + b_free)) # contracted axes to the front + a_fs = [ashape[i] for i in a_free] + b_fs = [bshape[i] for i in b_free] + m = int(np.prod(a_fs)) + k = int(np.prod([ashape[i] for i in a_axes])) + n = int(np.prod(b_fs)) + a3 = be.reshape(a_t, (m, k, 1)) + b3 = be.reshape(b_t, (1, k, n)) + return a3, b3, (tuple(a_fs) + tuple(b_fs)) + + +def _tropical_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + """max-plus tensordot: Y[i,j] = max_k (A[i,k] + B[k,j]).""" + a3, b3, out_shape = _pair_layout(be, a, b, axes) + s = a3 + b3 # native broadcast add -> (m, k, n) + red = be.max(s, axis=1) # max over contracted axis (tropical addition) + return be.reshape(red, out_shape) + + +def _expand_to_layout( + be: Backend, t: Tensor, idxs: Sequence[str], full: Sequence[str] +) -> Tensor: + """Reshape/transpose ``t`` (axes == ``idxs``) into the ``full`` index layout, + inserting size-1 axes for indices not in ``idxs`` (enables broadcasting).""" + shape = tuple(int(x) for x in be.shape_tuple(t)) + present = {c: shape[i] for i, c in enumerate(idxs)} + sub = [c for c in full if c in present] # present indices in full order + tt = be.transpose(t, tuple(idxs.index(c) for c in sub)) + newshape = tuple(present[c] if c in present else 1 for c in full) + return be.reshape(tt, newshape) + + +def _tropical_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: + """max-plus einsum: product -> +, sum -> max. Handles 1- or 2-operand forms.""" + if len(operands) == 1: + a = operands[0] + in_str, _sep, out_str = eq.partition("->") + lhs = in_str.split(",")[0] + rhs = out_str + if len(set(lhs)) != len(lhs): + resolved = "".join(dict.fromkeys(lhs)) + a = be.einsum(lhs + "->" + resolved, a) + lhs = resolved + contract = [c for c in lhs if c not in rhs] + for ax in sorted([lhs.index(c) for c in contract], reverse=True): + a = be.max(a, axis=ax) + remaining = [c for c in lhs if c not in contract] + return be.transpose(a, tuple(remaining.index(c) for c in rhs)) + a, b = operands + lhs, rhs = eq.split("->") + ia_s, ib_s = lhs.split(",") + ia, ib = list(ia_s), list(ib_s) + if len(set(ia)) != len(ia): + resolved = "".join(dict.fromkeys(ia)) + a = be.einsum("".join(ia) + "->" + resolved, a) + ia = list(resolved) + if len(set(ib)) != len(ib): + resolved = "".join(dict.fromkeys(ib)) + b = be.einsum("".join(ib) + "->" + resolved, b) + ib = list(resolved) + all_idx = list(dict.fromkeys(ia + ib)) + out_idx = list(rhs) + contract = [c for c in all_idx if c not in out_idx] + s = _expand_to_layout(be, a, ia, all_idx) + _expand_to_layout(be, b, ib, all_idx) + for ax in sorted([all_idx.index(c) for c in contract], reverse=True): + s = be.max(s, axis=ax) + remaining = [c for c in all_idx if c not in contract] + return be.transpose(s, tuple(remaining.index(c) for c in out_idx)) + + +class MaxPlusAlgebra(ContractionAlgebra): + """Tropical (max, +) semiring: addition -> max, multiplication -> +.""" + + name = "maxplus" + # representation defaults to IdentityRepresentation (real tensors, no codec): + # leaves enter/exit the contraction unchanged; only the kernels differ. + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + return _tropical_tensordot(be, a, b, axes) + + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + return _tropical_einsum(be, eq, *operands) + + +# ===== Section 2: counting (energy, degeneracy) ===== + + +def split_energy_count(stacked: Tensor) -> "tuple[Tensor, Tensor]": + """Split a stacked ``[..., 2]`` tensor into ``(energy, count)`` numpy arrays. + + The last axis is interpreted as ``[..., 0] = energy, [..., 1] = count``. + """ + arr = stacked if isinstance(stacked, np.ndarray) else np.asarray(stacked) + return arr[..., 0], arr[..., 1] + + +def _stack_last(be: Backend, x: Tensor, n: Tensor) -> Tensor: + """Stack two same-shape tensors along a new trailing axis (portable). + + All tc-ng concrete backends (numpy, jax, torch, tensorflow, cupy) implement + ``be.stack`` via ``tensorcircuit.backends.abstract_backend``, so the portable + path is used directly. ``be.stack`` takes a Python sequence and inserts a new + axis; ``axis=-1`` puts it last to match the ``[..., 2]`` convention. + """ + return be.stack([x, n], axis=-1) + + +def _counting_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + """(max-plus energy, degeneracy count) pairwise contraction. + + ``a`` and ``b`` are stacked ``[..., 2]`` tensors (``[..., 0]`` = energy, + ``[..., 1]`` = count). The energy stream follows max-plus; the count stream + sums ``a_count * b_count`` over the contracted positions that achieve the + energy max (within ``_EPS``), so each max-tie contributes its multiplicity. + """ + a3x, b3x, out_shape = _pair_layout(be, a[..., 0], b[..., 0], axes) + s = a3x + b3x # (m, k, n) energy pair sums + y = be.max(s, axis=1) # (m, n) max energy per output slot + # abs(s - y) < _EPS, broadcast across the contracted axis k: reshape y to (m, 1, n) + m_n = tuple(int(d) for d in be.shape_tuple(y)) + y_b = be.reshape(y, (m_n[0], 1, m_n[1])) + mask = be.abs(s - y_b) < _EPS + a3n, b3n, _ = _pair_layout(be, a[..., 1], b[..., 1], axes) + pn = a3n * b3n # (m, k, n) pairwise count products + cy = be.reshape(be.sum(pn * mask, axis=1), out_shape) + y_shaped = be.reshape(y, out_shape) + return _stack_last(be, y_shaped, cy) + + +def _insert_axis_of( + be: Backend, reduced: Tensor, axis: int, full_shape: Sequence[int] +) -> Tensor: + """Reshape an axis-reduced tensor so its ``axis`` is size 1, restoring the + pre-reduce rank for broadcasting against the original tensor. + + ``be.max`` lacks ``keepdims``; this reshape re-inserts the size-1 axis so + ``reduced`` broadcasts against a tensor of shape ``full_shape``. + """ + new_shape = tuple( + 1 if i == axis else int(full_shape[i]) for i in range(len(full_shape)) + ) + return be.reshape(reduced, new_shape) + + +def _resolve_repeats( + be: Backend, pair: Tensor, idxs: Sequence[str] +) -> "tuple[Tensor, Tensor, list[str]]": + """Resolve intra-operand repeated indices via per-stream diagonal gather. + + Pure per-stream indexing -- the energy and count streams are gathered + independently (``pair[..., 0]`` and ``pair[..., 1]``) so the trailing + ``[..., 2]`` stack axis is never treated as an index by ``be.einsum``. The + max/tie degeneracy logic lives entirely in the later reduction; this helper + only rewrites each stream's index layout. + + Returns ``(energy, count, resolved_idxs)`` where ``resolved_idxs`` has no + repeats. If ``idxs`` has no repeats, the streams are sliced out unchanged + and ``idxs`` is returned as-is (passthrough -- 2-operand behavior is + identical to the pre-helper code path). + """ + if len(set(idxs)) != len(idxs): + resolved = "".join(dict.fromkeys(idxs)) + e = be.einsum("".join(idxs) + "->" + resolved, pair[..., 0]) + n = be.einsum("".join(idxs) + "->" + resolved, pair[..., 1]) + return e, n, list(resolved) + return pair[..., 0], pair[..., 1], list(idxs) + + +def _counting_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: + """(energy, degeneracy) einsum over max-plus. + + Handles the 1-operand and 2-operand forms that cotengra can emit. + + Mirrors ``_tropical_einsum``: each operand is broadcast to the full index + layout via ``_expand_to_layout``; the energy stream sums then takes the max + over each contracted axis; the count stream sums ``a_count * b_count`` only + at positions within ``_EPS`` of the running energy max (ties contribute). + + Intra-operand repeated indices (diagonal/trace) are resolved per-stream via + ``_resolve_repeats`` (pure diagonal gather -- no max/tie logic, which lives + in the reduction). This is the per-stream analogue of ``_tropical_einsum``'s + ``be.einsum`` gather, split across the two trailing-axis streams so the + ``[..., 2]`` stack axis is preserved rather than treated as an index. + """ + if len(operands) == 1: + a = operands[0] + in_str, _sep, out_str = eq.partition("->") + lhs = in_str.split(",")[0] + rhs = out_str + e, n, idxs = _resolve_repeats(be, a, list(lhs)) + contract = [c for c in idxs if c not in rhs] + for ax in sorted([idxs.index(c) for c in contract], reverse=True): + shp = tuple(int(d) for d in be.shape_tuple(e)) + y_ax = be.max(e, axis=ax) + y_b = _insert_axis_of(be, y_ax, ax, shp) + mask = be.abs(e - y_b) < _EPS + e = y_ax + n = be.sum(n * mask, axis=ax) + remaining = [c for c in idxs if c not in contract] + out_e = be.transpose(e, tuple(remaining.index(c) for c in rhs)) + out_n = be.transpose(n, tuple(remaining.index(c) for c in rhs)) + return _stack_last(be, out_e, out_n) + a, b = operands + lhs, rhs = eq.split("->") + ia_s, ib_s = lhs.split(",") + a_e, a_n, ia = _resolve_repeats(be, a, list(ia_s)) + b_e, b_n, ib = _resolve_repeats(be, b, list(ib_s)) + all_idx = list(dict.fromkeys(ia + ib)) + out_idx = list(rhs) + contract = [c for c in all_idx if c not in out_idx] + # Energy pair-sum and count product broadcast over the full index layout. + sx = _expand_to_layout(be, a_e, ia, all_idx) + _expand_to_layout( + be, b_e, ib, all_idx + ) + pn = _expand_to_layout(be, a_n, ia, all_idx) * _expand_to_layout( + be, b_n, ib, all_idx + ) + # Reduce over contracted axes (highest index first so earlier indices stay valid). + for ax in sorted([all_idx.index(c) for c in contract], reverse=True): + sx_shape = tuple(int(d) for d in be.shape_tuple(sx)) + y_ax = be.max(sx, axis=ax) # axis removed + y_b = _insert_axis_of(be, y_ax, ax, sx_shape) # axis size 1 -> broadcasts + mask = be.abs(sx - y_b) < _EPS + sx = y_ax + pn = be.sum(pn * mask, axis=ax) + remaining = [c for c in all_idx if c not in contract] + out_x = be.transpose(sx, tuple(remaining.index(c) for c in out_idx)) + out_n = be.transpose(pn, tuple(remaining.index(c) for c in out_idx)) + return _stack_last(be, out_x, out_n) + + +class CountingRepresentation(Representation): + """Attach count=1 (the counting-semiring multiplicative identity) to each + leaf; decode splits the final stacked ``[..., 2]`` tensor into energy + (primary) + count (aux). + + encode: ``t -> stack([t, ones_like(t, float64)], axis=-1)``. Per-tensor, + topology-agnostic, run once on the leaves before any pairwise contraction. + decode: ``tensor[..., 0]`` is the energy (primary, rank == len(output_set)); + ``tensor[..., 1]`` is the degeneracy count (aux, co-indexed with energy). + """ + + name = "counting" + + def encode(self, be: Backend, tensors: List[Tensor]) -> List[Tensor]: + out = [] + for t in tensors: + ones = be.ones_like(t, dtype="float64") + out.append(be.stack([t, ones], axis=-1)) + return out + + def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]: + return tensor[..., 0], {"count": tensor[..., 1]} + + +class CountingTropicalAlgebra(ContractionAlgebra): + """Counting tropical algebra: (energy, degeneracy) over max-plus. + + Carries ``CountingRepresentation``: encode attaches count=1 to each leaf; + decode splits the final pair into energy (primary) + count (aux, stashed + via ``cons._stash_aux_outputs`` for ``degeneracy()`` to read). No + ``on_contraction_start`` hook is needed: the unconditional aux clear in + ``cons._algebraic_base_contraction`` already wipes stale state per + contraction (standard or non-standard). + """ + + name = "counting_maxplus" + representation = CountingRepresentation() + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + return _counting_tensordot(be, a, b, axes) + + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + return _counting_einsum(be, eq, *operands) + + +def degeneracy() -> Optional[Any]: + """Degeneracy (number of optimal configs) of the most recent counting + contraction. + + Call inside the ``counting_tropical()`` block (like + ``recover_configuration``) after a contraction has run. Returns the count + tensor (same shape as the primary energy) stashed by + ``CountingRepresentation.decode``; returns ``None`` if no counting decode + has stashed a count for the current contraction. + + The aux side-channel is cleared at the start of every contraction that + routes through ``cons._algebraic_base_contraction`` (the unconditional + ``_stash_aux_outputs({})`` at the top of that function), so a counting + contraction followed by another algebraic contraction wipes stale state. + Note that ``cons._base`` only routes to ``_algebraic_base_contraction`` + for non-standard algebras, hyperedge inputs, or ``use_primitives=True``; + a standard contraction on the original ``_base`` path does NOT clear aux. + """ + return cons._aux_outputs().get("count") + + +# ===== Section 3: tracking + configuration recovery ===== + +# Per-call argmax records, appended in call-order (== tree.traverse() order). +# Reset by on_contraction_start (fired by cons._algebraic_base_contraction) at +# the start of each tracked contraction so it always reflects exactly the most +# recent contraction. +_trace: List[Dict[str, Any]] = [] + +# Context stashed by the tracking hooks (fired by cons._algebraic_base_contraction): +# the cotengra tree plus the extracted topology (raw leaf arrays + index terms). +_ctx: Dict[str, Any] = { + "tree": None, + "input_sets": None, + "raw_tensors": None, +} + + +def _reset_trace() -> None: + """Clear the argmax trace (called by on_contraction_start, fired by cons._algebraic_base_contraction). + + Also drops any stashed tree so a failed/short-circuited contraction cannot + be confused with the previous one by ``recover_configuration``. + """ + _trace.clear() + _ctx["tree"] = None + + +def _set_tracking_context( + tree: Any, + input_sets: Optional[Sequence[Sequence[Any]]] = None, + raw_tensors: Optional[Sequence[Any]] = None, +) -> None: + """Stash the cotengra tree (and optionally the leaf topology) for backtracking.""" + _ctx["tree"] = tree + _ctx["input_sets"] = input_sets + _ctx["raw_tensors"] = raw_tensors + + +def get_recorded_topology() -> Tuple[Any, Any, Any]: + """Return ``(tree, input_sets, raw_tensors)`` stashed by the last tracked + contraction. The test harness uses this to verify a recovered config's + energy without needing to know cotengra's internal symbol->spin mapping.""" + return _ctx["tree"], _ctx["input_sets"], _ctx["raw_tensors"] + + +def _axes_to_lists(ashape: Sequence[int], axes: Any) -> Tuple[List[int], List[int]]: + """Normalize a tensordot ``axes`` arg to ``(a_axes, b_axes)`` (mirror of + ``_pair_layout``).""" + if isinstance(axes, int): + a_axes = list(range(len(ashape) - axes, len(ashape))) + b_axes = list(range(0, axes)) + else: + a_axes, b_axes = list(axes[0]), list(axes[1]) + return a_axes, b_axes + + +def _tracking_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + """max-plus tensordot (identical result to ``_tropical_tensordot``) that + additionally records the argmax over the contracted axis per output pos. + + Record shape: ``argmax`` has the pairwise output shape in the algebra's + natural order (``a_free + b_free``); values are the *flattened* contracted + index in row-major over ``a_axes`` (so ``np.unravel_index`` with + ``contract_dims`` recovers per-label values). ``contract_dims`` follows the + order of ``a_axes`` (== ``tree.get_tensordot_axes(p)[0]``). + + Argmax uses ``be.argmax`` (portable across numpy/jax/tensorflow); the + result is materialised to numpy for the backtracking index walk. + """ + a3, b3, out_shape = _pair_layout(be, a, b, axes) # (m,k,1),(1,k,n) + s = a3 + b3 # (m,k,n) + red = be.max(s, axis=1) # (m,n) max-plus reduction + out = be.reshape(red, out_shape) + + ashape = tuple(int(x) for x in be.shape_tuple(a)) + a_axes, _b_axes = _axes_to_lists(ashape, axes) + contract_dims = tuple(ashape[ax] for ax in a_axes) + + am = be.argmax(s, axis=1) # (m,n) -> flattened contracted index + am = be.reshape(am, out_shape) + _trace.append( + { + "kind": "td", + "argmax": np.asarray(am), + "contract_dims": contract_dims, + } + ) + return out + + +def _tracking_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: + """max-plus einsum (identical result to ``_tropical_einsum``) that records + the argmax over the contracted index subspace per output position. + + For 2-operand forms the pair-sum is built over the full index layout (as in + ``_tropical_einsum``), transposed to ``(out_labels + contract_labels)``, + the contracted axes flattened into one trailing axis, and ``be.argmax`` + taken over it -- vectorised and portable (numpy/jax/tensorflow). Single- + operand forms record a no-op entry (no pairwise choice to backtrack). + """ + if len(operands) == 1: + _trace.append( + { + "kind": "ein1", + "argmax": None, + "contract_dims": (), + "contract_labels": (), + "eq": eq, + } + ) + else: + a, b = operands + lhs, rhs = eq.split("->") + ia_s, ib_s = lhs.split(",") + ia, ib = list(ia_s), list(ib_s) + out_labels = list(rhs) + all_idx = list(dict.fromkeys(ia + ib)) + contract_labels = [c for c in all_idx if c not in out_labels] + + # Full-layout pair sum (same construction as _tropical_einsum). + s = _expand_to_layout_pair(be, a, b, ia, ib, all_idx) + + if not contract_labels: + # Pure outer product (hyperedge with no contraction): no argmax to take. + _trace.append( + { + "kind": "ein2", + "argmax": None, + "contract_dims": (), + "contract_labels": (), + "eq": eq, + } + ) + else: + # Reorder to (out_labels + contract_labels) and flatten the contracted axes. + perm = tuple(all_idx.index(c) for c in out_labels + contract_labels) + s = be.transpose(s, perm) + full_shape = tuple(int(d) for d in be.shape_tuple(s)) + n_out = len(out_labels) + out_shape = full_shape[:n_out] + contract_dims = full_shape[n_out:] + flat = int(np.prod(contract_dims)) + s_flat = be.reshape(s, out_shape + (flat,)) + am = be.argmax( + s_flat, axis=-1 + ) # shape == out_shape, row-major over contract + _trace.append( + { + "kind": "ein2", + "argmax": np.asarray(am), + "contract_dims": tuple(contract_dims), + "contract_labels": tuple(contract_labels), + "eq": eq, + } + ) + return _tropical_einsum(be, eq, *operands) + + +def _expand_to_layout_pair( + be: Backend, + a: Tensor, + b: Tensor, + ia: Sequence[str], + ib: Sequence[str], + all_idx: Sequence[str], +) -> Tensor: + """Build the full-layout pair-sum ``a + b`` broadcast over ``all_idx``. + + Mirrors the pair-sum construction inside ``_tropical_einsum`` (factors + ``_expand_to_layout`` for both operands). Kept local so this module does not + reach into a private helper whose signature may change. + """ + + return _expand_to_layout(be, a, ia, all_idx) + _expand_to_layout(be, b, ib, all_idx) + + +class MaxPlusTrackingAlgebra(MaxPlusAlgebra): + """Max-plus algebra that records the per-step argmax for config recovery. + + Produces the same contraction value as ``MaxPlusAlgebra``; the only + side-effect is appending an argmax record to ``_trace`` per pairwise call. + The trace is reset by ``on_contraction_start`` (fired by + ``cons._algebraic_base_contraction``) at the start of each tracked + contraction, so ``recover_configuration`` always reflects the most recent + contraction. + """ + + name = "maxplus_tracking" + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + return _tracking_tensordot(be, a, b, axes) + + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + return _tracking_einsum(be, eq, *operands) + + def on_contraction_start(self, nodes: Any) -> None: + _reset_trace() + try: + _raw, _inputs, _output, _sizes = cons._extract_topology(nodes) + _set_tracking_context(tree=None, input_sets=_inputs, raw_tensors=_raw) + except Exception: + logger.debug( + "tracking topology stash failed; recover_configuration may raise " + "a trace/tree mismatch later", + exc_info=True, + ) + _set_tracking_context(tree=None) + + def on_contractor_ready(self, tree: Any) -> None: + _ctx["tree"] = tree + + +def _td_backtrack_step( + tree: Any, + p: Any, + l: Any, + r: Any, + rec: Dict[str, Any], + p_pos: Tuple[int, ...], + assignment: Dict[Any, int], +) -> Tuple[Tuple[int, ...], Tuple[int, ...]]: + """Propagate one tensordot contraction's optimal position to its children. + + Returns ``(l_pos, r_pos)`` (positions in ``get_inds(l)`` / ``get_inds(r)`` + order). Side-effect: records the contracted index-label values into + ``assignment``. + + ``get_tensordot_perm(p)`` is the transpose cotengra applies AFTER the + algebra call, mapping the algebra's natural output order (``l_free + + r_free``) to ``get_inds(p)``. It is therefore INVERTED when indexing the + recorded argmax (which lives in the algebra's natural order) by the + canonical (``get_inds(p)``) position: ``td_pos[perm[i]] = p_pos[i]``. + """ + l_inds = list(tree.get_inds(l)) + r_inds = list(tree.get_inds(r)) + l_axes, r_axes = tree.get_tensordot_axes(p) + perm = tree.get_tensordot_perm(p) + + # Convert canonical (get_inds(p)) position -> algebra-natural (td) order. + if perm is None: + td_pos = list(p_pos) + else: + td_pos = [0] * len(p_pos) + for i, v in enumerate(p_pos): + td_pos[perm[i]] = v + td_pos_t = tuple(td_pos) + + argmax = rec["argmax"] + if argmax is None: + # No contracted axis (should not happen for a tensordot step); nothing + # to unravel -- pass the position through unchanged. + return tuple(p_pos[: len(l_inds)]), tuple(p_pos[len(l_inds) :]) + + k = int(argmax[td_pos_t]) + contract_dims = rec["contract_dims"] + if contract_dims: + unraveled = np.unravel_index(k, contract_dims) # per contracted axis + else: + unraveled = () + + # Contracted labels follow l_axes order (== the algebra's a_axes order). + contract_labels = [l_inds[ax] for ax in l_axes] + for i, lbl in enumerate(contract_labels): + assignment[lbl] = int(unraveled[i]) + + # Split td_pos into the (l_free, r_free) halves (td-natural order). + l_free_labels = [c for i, c in enumerate(l_inds) if i not in set(l_axes)] + r_free_labels = [c for i, c in enumerate(r_inds) if i not in set(r_axes)] + n_lf = len(l_free_labels) + l_free_vals = td_pos[:n_lf] + r_free_vals = td_pos[n_lf:] + + l_pos = [0] * len(l_inds) + for j, lbl in enumerate(l_free_labels): + l_pos[l_inds.index(lbl)] = l_free_vals[j] + for i, ax in enumerate(l_axes): + l_pos[ax] = int(unraveled[i]) + + r_pos = [0] * len(r_inds) + for j, lbl in enumerate(r_free_labels): + r_pos[r_inds.index(lbl)] = r_free_vals[j] + for i, ax in enumerate(r_axes): + r_pos[ax] = int(unraveled[i]) # r_axes[i] pairs with l_axes[i] (same label) + + return tuple(l_pos), tuple(r_pos) + + +def _ein_backtrack_step( + p_pos: Tuple[int, ...], + rec: Dict[str, Any], +) -> Dict[Any, int]: + """Recover the contracted-label assignment for one einsum step (given the + output position). Used when ``tree.get_can_dot(p)`` is False.""" + argmax = rec["argmax"] + contract_dims = rec["contract_dims"] + if argmax is None or not contract_dims: + return {} + k = int(argmax[p_pos]) + unraveled = np.unravel_index(k, contract_dims) + return {lbl: int(v) for lbl, v in zip(rec["contract_labels"], unraveled)} + + +def _validate_trace(tree: Any) -> None: + """Validate that a stashed contraction tree exists and that ``_trace`` + length matches its number of contractions; raise RuntimeError otherwise. + + Called at the top of ``recover_configuration`` so the trace/tree invariant + holds before backtracking begins. + """ + if tree is None: + raise RuntimeError( + "recover_configuration: no stashed tree -- contract under " + "tropical(track=True) first." + ) + n_nodes = len(list(tree.traverse())) + if len(_trace) != n_nodes: + raise RuntimeError( + f"recover_configuration: trace length ({len(_trace)}) != tree " + f"contractions ({n_nodes}); call after exactly one tracked " + "contraction." + ) + + +def _finalize_from_leaves( + tree: Any, + opt_pos: Dict[Any, Tuple[int, ...]], + assignment: Dict[Any, int], +) -> None: + """Finalize ``assignment`` from leaf positions, catching labels that appear + only on a leaf axis (never seen on a contracted/tensordot step). + Mutates ``assignment`` in place. + """ + for leaf in tree.gen_leaves(): + pos = opt_pos.get(leaf) + if pos is None: + continue + for ax, lbl in enumerate(list(tree.get_inds(leaf))): + assignment[lbl] = int(pos[ax]) + + +def recover_configuration() -> Dict[Any, int]: + """Walk the stashed tree top-down and recover each index label's optimal + value. Returns ``{index_label: value}``. + + Requires that the contraction was performed under ``MaxPlusTrackingAlgebra`` + (i.e. ``tropical(track=True)``) so that ``_trace`` and ``_ctx['tree']`` are + populated, and called after exactly one tracked contraction (the trace is + reset per contraction by ``on_contraction_start``). + + Only scalar roots are supported (a full contraction to an energy, the + Ising use case). If the contraction has dangling/free output indices + (``len(tree.output) != 0``) a ``NotImplementedError`` is raised: the + non-scalar backtracking path has known ordering bugs (``tree.output`` order + vs result-shape order; output-label values lost in the einsum branch) that + would produce wrong configs, so it is gated rather than shipping wrong answers. + Contract to a scalar first. + + Tie-breaking is first-argument-wins (lowest flattened contracted index on + ties), so the returned configuration is *an* optimum; degenerate optima are + not enumerated (that is Task B's remit). + """ + tree = _ctx["tree"] + _validate_trace(tree) + + trav = list(tree.traverse()) # call-order == algebra call order + node_rec: Dict[Any, Dict[str, Any]] = {} + for (p, _l, _r), rec in zip(trav, _trace): + node_rec[p] = rec + + # Scalar-only: gate the non-scalar (free/output-index) root, whose + # backtracking wiring has known ordering bugs (tree.output order vs + # result-shape order; output-label values lost in the einsum branch). The + # tested scalar path below stays intact; ``get_tensordot_perm`` inversion + # logic (correct) is still exercised by the scalar canary + synthetic test. + n_out = len(tree.output) + if n_out != 0: + raise NotImplementedError( + "non-scalar configuration recovery not supported; " + "contract to a scalar first" + ) + root_pos: Tuple[int, ...] = () + + assignment: Dict[Any, int] = {} + opt_pos: Dict[Any, Tuple[int, ...]] = {} + opt_pos[tree.root] = root_pos + + for p, l, r in tree.descend(): + rec = node_rec[p] + p_pos = opt_pos[p] + if rec["kind"] == "td": + l_pos, r_pos = _td_backtrack_step(tree, p, l, r, rec, p_pos, assignment) + opt_pos[l] = l_pos + opt_pos[r] = r_pos + else: + # einsum (hyperedge): recover contracted labels; the child + # positions in canonical order are derived from the labels. + cont = _ein_backtrack_step(p_pos, rec) + assignment.update(cont) + # Reconstruct child positions from the label->value map so deeper + # tensordot steps (which read opt_pos) stay consistent. Each child's + # surviving labels are exactly its get_inds restricted to known vals. + l_inds = list(tree.get_inds(l)) + r_inds = list(tree.get_inds(r)) + opt_pos[l] = tuple(assignment[c] if c in assignment else 0 for c in l_inds) + opt_pos[r] = tuple(assignment[c] if c in assignment else 0 for c in r_inds) + + _finalize_from_leaves(tree, opt_pos, assignment) + return assignment + + +# ===== Section 4: context managers ===== + + +@contextlib.contextmanager +def tropical(track: bool = False) -> Iterator[None]: + """Contract under the max-plus (tropical) algebra within the block. + + Swaps ``cons._contraction_algebra`` directly (no monkey-patch activation): + the in-source ``cons._base`` routes to ``_algebraic_base_contraction`` for + any non-standard algebra, so this is sufficient. + + ``track=True`` switches to ``MaxPlusTrackingAlgebra`` so + ``recover_configuration()`` can recover the optimal configuration after the + contraction. Off by default -> zero behaviour change relative to plain + max-plus. + """ + algebra = MaxPlusTrackingAlgebra() if track else MaxPlusAlgebra() + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(algebra) + try: + yield + finally: + cons.set_contraction_algebra(prev) + + +@contextlib.contextmanager +def counting_tropical() -> Iterator[None]: + """Contract under the counting (energy, degeneracy) max-plus algebra within + the block.""" + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(CountingTropicalAlgebra()) + try: + yield + finally: + cons.set_contraction_algebra(prev) diff --git a/examples/tropical_ising.py b/examples/tropical_ising.py new file mode 100644 index 00000000..9dbd1abb --- /dev/null +++ b/examples/tropical_ising.py @@ -0,0 +1,66 @@ +"""Tropical (max-plus) Ising ground-state energy via TensorCircuit-NG contraction. + +Run: PYTHONPATH= python examples/tropical_ising.py +""" + +import itertools +import numpy as np +import tensornetwork as tn +import tensorcircuit as tc +import tensorcircuit.cons as cons +from applications.tropical_algebra import tropical + + +def brute_force_ground_energy(n, edges, j_vals, h): + best = np.inf + for cfg in itertools.product([-1, 1], repeat=n): + e = -sum(jij * cfg[i] * cfg[j] for (i, j), jij in zip(edges, j_vals)) - sum( + hi * cfg[i] for i, hi in enumerate(h) + ) + best = min(best, e) + return best + + +def build_tn(n, edges, j_vals, h): + be = tc.backend + degree = [0] * n + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + nodes, copy_nodes = [], {} + for i in range(n): + cn = tn.CopyNode(degree[i] + 1, 2) # rank (degree+1) delta, dim 2 -> hyperedge + copy_nodes[i] = cn + nodes.append(cn) + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + leg = [1] * n + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg[i]]) + tn.connect(ten[1], copy_nodes[j][leg[j]]) + leg[i] += 1 + leg[j] += 1 + nodes.append(ten) + return nodes + + +def main(): + n = 5 + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + nodes = build_tn(n, edges, J, h) + with tropical(): + val = float(np.array(cons.contractor(nodes, output_edge_order=[]).tensor)) + e_ground = brute_force_ground_energy(n, edges, J, h) + print(f"tropical contraction = {val:.6f}") + print(f"-E_ground (brute) = {-e_ground:.6f}") + print(f"match: {np.isclose(val, -e_ground)}") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index c1984d1b..0a2a7db0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ filterwarnings = [ "ignore::DeprecationWarning", "ignore:Explicitly requested dtype*:UserWarning", ] +pythonpath = ["."] [tool.mypy] diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index a4d5d865..1b6edaff 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -3,7 +3,7 @@ pytest==7.4.4 pytest-cov pytest-benchmark pytest-xdist -black[jupyter] +black[jupyter]==26.5.1 sphinx>=4.0 pytest-lazy-fixture pylint==3.2.6 diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 6cb13b77..7d591659 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -4,6 +4,7 @@ # pylint: disable=invalid-name +import functools import logging import sys import time @@ -138,6 +139,43 @@ def set_tensornetwork_backend( set_tensornetwork_backend() +# --- ContractionAlgebra state (mirrors dtypestr/set_dtype pattern) --- +from .contraction_algebra import StandardAlgebra as _StandardAlgebra +from .contraction_algebra.base import ( + ContractionAlgebra as _ContractionAlgebra, +) + +_contraction_algebra: _ContractionAlgebra = _StandardAlgebra() + + +def get_contraction_algebra() -> _ContractionAlgebra: + return _contraction_algebra + + +def set_contraction_algebra(alg: _ContractionAlgebra) -> None: + global _contraction_algebra + _contraction_algebra = alg + + +# Aux output side-channel for multi-output (non-standard) algebras. Cleared and +# refilled on each non-standard contraction by ``_stash_aux_outputs``; the +# primary tensor is returned through the normal contraction return value, while +# aux (e.g. counting degeneracy, sharing the primary's physical axis order) is +# stashed here for the caller to read via ``_aux_outputs``. +_aux_outputs_store: Dict[str, Any] = {} + + +def _stash_aux_outputs(aux: Dict[str, Any]) -> None: + """Stash the most recent contraction's aux outputs (side channel for + multi-output algebras, e.g. counting degeneracy). Cleared per contraction.""" + _aux_outputs_store.clear() + _aux_outputs_store.update(aux) + + +def _aux_outputs() -> Dict[str, Any]: + """Test/accessor: snapshot of the stashed aux outputs.""" + return dict(_aux_outputs_store) + def set_function_backend(backend: Optional[str] = None) -> Callable[..., Any]: """ @@ -297,6 +335,12 @@ def _sizen(node: tn.Node, is_log: bool = False) -> int: def _merge_single_gates( nodes: List[Any], total_size: Optional[int] = None ) -> Tuple[List[Any], int]: + if not isinstance(_contraction_algebra, _StandardAlgebra): + # _merge_single_gates contracts via tn.contract_parallel (native sum-product), + # bypassing the algebra kernel — skip it under a non-standard algebra. + if total_size is None: + total_size = sum([_sizen(t) for t in nodes]) + return nodes, total_size # TODO(@refraction-ray): investigate whether too much copy here so that staging is slow for large circuit nodes = list(nodes) if total_size is None: @@ -697,48 +741,74 @@ def _wrap_omeco_optimizer(optimizer: Any) -> Any: return optimizer -def _algebraic_base_contraction( - nodes: List[tn.Node], - algorithm: Any, - output_edge_order: Optional[Sequence[tn.Edge]] = None, - ignore_edge_order: bool = False, - **kws: Any, -) -> Any: - """ - Execute contraction using cotengra and autoray for bare tensors. - """ - import cotengra as ctg +def _stash_permuted_aux( + aux: Any, output_edge_order: Any, dangling_edges: Any, kbe: Any +) -> None: + """Apply the output_edge_order permutation to aux (count co-indexed with energy).""" + order = output_edge_order if output_edge_order is not None else dangling_edges + perm = [dangling_edges.index(e) for e in order] + _stash_aux_outputs({k: kbe.transpose(v, tuple(perm)) for k, v in aux.items()}) - raw_tensors, input_sets, output_set, size_dict = _extract_topology(nodes) - # Use the backend of the first node - be = nodes[0].backend +def _run_contraction( + raw_tensors: Any, + input_sets: Any, + output_set: Any, + size_dict: Any, + algorithm: Any, + ns: bool, + alg: Any, + kbe: Any, + be: Any, + strip_exponent: bool, + ctg: Any, +) -> Any: + """Run the (possibly single-tensor) contraction, returning (final, exponent).""" + exponent = 0.0 if len(raw_tensors) == 1: # Avoid cotengra bug for empty contraction paths - final_raw_tensor = be.einsum(input_sets[0] + "->" + output_set, *raw_tensors) - exponent = 0.0 + eq = input_sets[0] + "->" + output_set + final = alg.einsum(kbe, eq, *raw_tensors) if ns else be.einsum(eq, *raw_tensors) else: path = algorithm(input_sets, output_set, size_dict) logger.info("the contraction path is given as %s" % str(path)) - tree = ctg.ContractionTree.from_path( input_sets, output_set, size_dict, path=path ) - - # Use autoray to keep AD and JIT support across backends - # Note: cotengra's make_contractor handles the orchestration - if not kws.get("strip_exponent", False): + if ns: + alg.on_contractor_ready(tree) + impl = ( + functools.partial(alg.einsum, kbe), + functools.partial(alg.tensordot, kbe), + ) + contractor = ctg.core.make_contractor( + tree, implementation=impl, prefer_einsum=alg.prefer_einsum + ) + final = contractor(*raw_tensors) + elif not strip_exponent: + # autoray keeps AD and JIT support across backends contractor = ctg.core.make_contractor(tree, implementation="autoray") - final_raw_tensor = contractor(*raw_tensors) + final = contractor(*raw_tensors) else: - final_raw_tensor, exponent = tree.contract(raw_tensors, strip_exponent=True) + final, exponent = tree.contract(raw_tensors, strip_exponent=True) + return final, exponent + + +def _decode_aux(ns: bool, rep: Any, kbe: Any, final: Any, output_set: Any) -> Any: + """Decode the contraction output under a non-standard algebra; else no aux.""" + if not ns: + return final, {} + primary, aux = rep.decode(kbe, final) + assert primary.ndim == len(output_set), ( + "representation.decode primary rank %d != len(output_set) %d; " + "decode must strip non-physical storage axes before tn.Node wraps it" + % (primary.ndim, len(output_set)) + ) + return primary, aux - final_node = tn.Node(final_raw_tensor, backend=be) - # Resolve dangling edges in the same order as in _extract_topology - dangling_edges = sorted_edges(tn.get_subgraph_dangling(nodes)) - - # Update the edges to point to the new final_node +def _rewire_dangling_edges(final_node: Any, dangling_edges: Any, nodes: Any) -> None: + """Point every dangling edge at final_node, preserving topology order.""" for i, edge in enumerate(dangling_edges): if edge.node1 in nodes: edge.node1 = final_node @@ -748,12 +818,73 @@ def _algebraic_base_contraction( edge.axis2 = i final_node.edges = list(dangling_edges) + +def _algebraic_base_contraction( + nodes: List[tn.Node], + algorithm: Any, + output_edge_order: Optional[Sequence[tn.Edge]] = None, + ignore_edge_order: bool = False, + **kws: Any, +) -> Any: + """ + Execute contraction using cotengra and autoray for bare tensors. + """ + import cotengra as ctg + + raw_tensors, input_sets, output_set, size_dict = _extract_topology(nodes) + be = nodes[0].backend # tn backend: standard native ops + tn.Node wrap + + alg = get_contraction_algebra() + rep = alg.representation + ns = not isinstance(alg, _StandardAlgebra) + kbe = backend if ns else be # algebra kernels need the tc backend (max/argmax/...) + + # Unconditional clear: every contraction (standard or non-standard) wipes the + # aux side-channel so a subsequent ``degeneracy()`` cannot read a stale count + # left by an earlier counting contraction. Non-standard decodes then refill it. + _stash_aux_outputs({}) + if ns: + if kws.get("strip_exponent", False): + raise ValueError( + "strip_exponent is incompatible with a non-standard ContractionAlgebra" + ) + alg.on_contraction_start(nodes) + raw_tensors = rep.encode(kbe, raw_tensors) + + final, exponent = _run_contraction( + raw_tensors, + input_sets, + output_set, + size_dict, + algorithm, + ns, + alg, + kbe, + be, + kws.get("strip_exponent", False), + ctg, + ) + + final, aux = _decode_aux(ns, rep, kbe, final, output_set) + + final_node = tn.Node(final, backend=be) + + # Resolve dangling edges in the same order as in _extract_topology + dangling_edges = sorted_edges(tn.get_subgraph_dangling(nodes)) + + # Update the edges to point to the new final_node + _rewire_dangling_edges(final_node, dangling_edges, nodes) + if not ignore_edge_order: if output_edge_order is None: output_edge_order = dangling_edges final_node.reorder_edges(list(output_edge_order)) - if kws.get("strip_exponent", False): + # Apply the same output_edge_order permutation to aux (count co-indexed with energy) + if ns and aux: + _stash_permuted_aux(aux, output_edge_order, dangling_edges, kbe) + + if kws.get("strip_exponent", False) and not ns: return final_node, exponent return final_node @@ -902,7 +1033,8 @@ def _base( # 1. Resolve topology and check for hyperedges has_hyperedges = any(isinstance(n, tn.CopyNode) for n in nodes) - if use_primitives is True or (use_primitives is None and has_hyperedges): + _ns_alg = not isinstance(_contraction_algebra, _StandardAlgebra) + if use_primitives is True or _ns_alg or (use_primitives is None and has_hyperedges): # ========================================== # NEW ALGEBRAIC EXECUTION PATH (Opt-in) # ========================================== @@ -1142,6 +1274,7 @@ def set_contractor( contraction_info: bool = False, debug_level: int = 0, use_primitives: Optional[bool] = None, + algebra: Optional[_ContractionAlgebra] = None, **kws: Any, ) -> Callable[..., Any]: """ @@ -1162,6 +1295,20 @@ def set_contractor( :return: The new tensornetwork with its contractor set. :rtype: tn.Node """ + if algebra is not None: + if not isinstance(algebra, _StandardAlgebra): + use_primitives = True + if kws.get("preprocessing", False): + raise ValueError( + "preprocessing is incompatible with a non-standard " + "ContractionAlgebra (it contracts via native sum-product)" + ) + if kws.get("strip_exponent", False): + raise ValueError( + "strip_exponent is incompatible with a non-standard " + "ContractionAlgebra (its log-scaling assumes sum-product)" + ) + set_contraction_algebra(algebra) if not method: method = "greedy" # auto for small size fallbacks to dp, which has bug for now @@ -1291,11 +1438,13 @@ def wrapper(f: Callable[..., Any]) -> Callable[..., Any]: @wraps(f) def newf(*args: Any, **kws: Any) -> Any: old_contractor = getattr(thismodule, "contractor") + old_algebra = get_contraction_algebra() set_contractor(*confargs, **confkws) try: return f(*args, **kws) finally: _set_global_contractor(old_contractor) + set_contraction_algebra(old_algebra) return newf @@ -1311,11 +1460,13 @@ def runtime_contractor(*confargs: Any, **confkws: Any) -> Iterator[Any]: :rtype: Iterator[Any] """ old_contractor = getattr(thismodule, "contractor") + old_algebra = get_contraction_algebra() nc = set_contractor(*confargs, **confkws) try: yield nc finally: _set_global_contractor(old_contractor) + set_contraction_algebra(old_algebra) def split_rules( diff --git a/tensorcircuit/contraction_algebra/__init__.py b/tensorcircuit/contraction_algebra/__init__.py new file mode 100644 index 00000000..d831e06a --- /dev/null +++ b/tensorcircuit/contraction_algebra/__init__.py @@ -0,0 +1,27 @@ +"""ContractionAlgebra: a generic interface for swapping contraction primitives + +boundary representation, consulted by ``cons._algebraic_base_contraction``. + +Implement ``ContractionAlgebra`` (with a ``Representation``) and activate via +``cons.set_contraction_algebra(...)`` (or ``cons.set_contractor(algebra=...)``, +``cons.runtime_contractor(..., algebra=...)``). The in-source ``cons._base`` +routes any non-standard algebra to ``_algebraic_base_contraction``, which runs +encode -> algebra kernels -> decode; no monkey-patching is required. + +Reference applications: ``applications/tropical_algebra.py`` +(max-plus / counting / tracking) and ``applications/bcomplex32_algebra.py`` +(bf16 pair). +""" + +from .base import ( + ContractionAlgebra, + StandardAlgebra, + Representation, + IdentityRepresentation, +) + +__all__ = [ + "ContractionAlgebra", + "StandardAlgebra", + "Representation", + "IdentityRepresentation", +] diff --git a/tensorcircuit/contraction_algebra/base.py b/tensorcircuit/contraction_algebra/base.py new file mode 100644 index 00000000..4e680c14 --- /dev/null +++ b/tensorcircuit/contraction_algebra/base.py @@ -0,0 +1,95 @@ +"""ContractionAlgebra: generic interface for swapping contraction primitives + +boundary representation, consulted by ``cons._algebraic_base_contraction``. + +Two ABCs: +- ``ContractionAlgebra``: the arithmetic (tensordot/einsum kernels) + observation + hooks. Carries a ``Representation`` (boundary codec), default identity. +- ``Representation``: boundary encode (leaves -> physical storage) / decode + (final -> primary tensor + aux). Covers storage-split (complex<->pair) and + precision cast (f64<->bf16). Default identity. + +kernel must be closed over whatever its Representation produces (author's +contract; Representation is bundled in the algebra so users cannot mis-pair). +""" + +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Tuple + +Tensor = Any +Backend = Any + + +class Representation(ABC): + """Boundary codec between logical elements and physical storage.""" + + name: str = "abstract" + + @abstractmethod + def encode(self, be: Backend, tensors: List[Tensor]) -> List[Tensor]: + """Transform the leaf raw_tensors (one pass, per-tensor, topology-agnostic).""" + + @abstractmethod + def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]: + """Transform the final contracted tensor into ``(primary, aux)``. + ``primary`` must have rank == len(output_set) so ``tn.Node`` wraps it + consistently. ``aux`` carries side outputs (e.g., degeneracy), sharing + the primary's physical axis order (sliced from the same tensor).""" + + +class IdentityRepresentation(Representation): + """No-op codec: standard storage (real/complex scalars), no aux.""" + + name = "identity" + + def encode(self, be: Backend, tensors: List[Tensor]) -> List[Tensor]: + return tensors + + def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]: + return tensor, {} + + +class ContractionAlgebra(ABC): + """How to perform the two atomic ops of a contraction, plus observation hooks. + + cotengra decomposes any contraction into pairwise steps; each is a + ``tensordot`` (ordinary pair) or ``einsum`` (hyperedge / copy-node). An + algebra supplies both. The kernel must be closed over the layout its + ``representation`` produces. + """ + + name: str = "abstract" + representation: Representation = IdentityRepresentation() + # When operands carry a trailing non-physical storage axis (e.g. the + # complex pair), cotengra's tensordot mode post-transposes results via + # autoray and mishandles that extra axis (ValueError: axes don't match array) + # -- set True to force einsum-only execution, which forwards operands verbatim + # to ``einsum`` and skips the autoray transpose. Default False keeps tensordot + # mode: tropical config-recovery backtracking depends on the tensordot + # intermediate layout and would break under forced einsum. + prefer_einsum: bool = False + + @abstractmethod + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + """Pairwise contraction over ``axes`` (np.tensordot convention).""" + + @abstractmethod + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + """Pairwise einsum (cotengra may also call with 1 operand).""" + + def on_contraction_start(self, nodes: Any) -> None: + """Called before each non-standard contraction (default no-op).""" + + def on_contractor_ready(self, tree: Any) -> None: + """Called when the cotengra tree is built (default no-op).""" + + +class StandardAlgebra(ContractionAlgebra): + """The usual (sum, product) ring — identical to native backend behaviour.""" + + name = "standard" + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + return be.tensordot(a, b, axes) + + def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: + return be.einsum(eq, *operands) diff --git a/tests/_tropical_test_utils.py b/tests/_tropical_test_utils.py new file mode 100644 index 00000000..9620e73f --- /dev/null +++ b/tests/_tropical_test_utils.py @@ -0,0 +1,379 @@ +"""Shared helpers for the tropical Ising tensor-network tests. + +Extracted from ``tests/test_tropical_ising.py`` and +``tests/test_tropical_config.py`` to remove duplication between their +near-identical CopyNode Ising TN builders and brute-force energy +enumeration loops. +""" + +import itertools + +import numpy as np +import tensornetwork as tn + +import tensorcircuit as tc + + +def build_ising_tn(spins, edges, j_vals, h): + """Build the tropical (max-plus) Ising tensor network as tensornetwork nodes. + + Per spin: a ``tn.CopyNode`` of rank (degree+1), dimension 2. A CopyNode is a true + hyperedge hub -- tensornetwork/cotengra identify all its legs as one shared index + (the delta/copy constraint is structural, not stored as data). This is what makes + the contraction route hyperedge-containing pairs through the tropical **einsum** + branch and ordinary pairs through the tropical **tensordot** branch, covering both. + Per edge (i,j): Te[si,sj] = J*si*sj -> [[J,-J],[-J,J]]. + Per spin: field Tv[s] = h*s -> [h,-h]. + Contracting over max-plus yields max_cfg(-E) = -E_ground. + + Note: an earlier draft built the delta as a dense ``tn.Node`` (0 on diagonal, a NEG + proxy off-diagonal). That is a *regular* node, not a hyperedge, so it does not + trigger the einsum branch and -- worse -- ``cons.contractor``'s ``_merge_single_gates`` + preprocessing collapses the whole network to one node before the algebraic path + runs, silently bypassing the tropical primitives (result 0.0). Using ``tn.CopyNode`` + is both the faithful "copy node -> hyperedge" construction and the one that works + end-to-end through the public ``cons.contractor`` API. + """ + be = tc.backend + degree = dict.fromkeys(spins, 0) + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + + nodes = [] + copy_nodes = {} + # one field leg + one leg per incident edge + for i in spins: + cn = tn.CopyNode(degree[i] + 1, 2) # rank (degree+1) delta, dim 2 -> hyperedge + copy_nodes[i] = cn + nodes.append(cn) + # field tensor on the first leg + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + # bond tensors, each consuming one copy leg per endpoint + leg_idx = dict.fromkeys(spins, 1) # leg 0 reserved for field + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) + tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) + leg_idx[i] += 1 + leg_idx[j] += 1 + nodes.append(ten) + return nodes + + +def brute_force_energy(spins, edges, j_vals, h): + """Brute-force min energy: E = -sum J s_i s_j - sum h s_i.""" + n = len(spins) + best = np.inf + for cfg in itertools.product([-1, 1], repeat=n): + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + e -= jij * cfg[i] * cfg[j] + for i, hi in zip(spins, h): + e -= hi * cfg[i] + best = min(best, e) + return best + + +def brute_force_energy_and_degeneracy(n, edges, j_vals, h): + """Brute-force ground energy + degeneracy over spin configs. + + cfg in {0,1}^n; spin s = 1 - 2*cfg (cfg=0 -> s=+1). + E = -sum J s_i s_j - sum h s_i. Returns (best_e, degeneracy). + """ + best_e = None + deg = 0 + for cfg in itertools.product([0, 1], repeat=n): + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + si, sj = 1 - 2 * cfg[i], 1 - 2 * cfg[j] + e -= jij * si * sj + for i, hi in zip(range(n), h): + e -= hi * (1 - 2 * cfg[i]) + if best_e is None or e < best_e - 1e-9: + best_e, deg = e, 1 + elif abs(e - best_e) < 1e-9: + deg += 1 + return best_e, deg + + +# --- Task 10: non-scalar (one free spin) builder + brute force --- + + +# Topology shared by build_ring_with_free_spin / brute_nonscalar_counting. +# A 4-spin ring with spin 0 free: small enough to brute-force exactly, large +# enough to need both tensordot and einsum steps under cotengra. Integer +# couplings/fields keep the brute force exact (no rtol needed). +_RING_SPINS = (0, 1, 2, 3) +_RING_EDGES = ((0, 1), (1, 2), (2, 3), (3, 0)) +_RING_J = (-2, 1, 1, -2) # mixed ferro/antiferro of both magnitudes -> exercises +# # both tensordot and einsum branches under cotengra. +_RING_H = (2, 0, 2, -1) # tuned (via brute-force search) so the two free-spin +# # values give DIFFERENT energy AND DIFFERENT degeneracy: +# # v_free=0 -> E=7, N=2; v_free=1 -> E=5, N=1. A swapped +# # or permuted aux cannot pass the per-output assertion. +_FREE_SPIN = 0 + + +def build_ring_with_free_spin( + spins=_RING_SPINS, + edges=_RING_EDGES, + j_vals=_RING_J, + h=_RING_H, + free_spin=_FREE_SPIN, +): + """Build a small ring Ising with ONE spin's CopyNode given an extra dangling + leg (the free output index). Same ``Tv``/``Te``/``tn.CopyNode`` construction + as ``build_ising_tn`` -- the only difference is the free spin's CopyNode has + one ADDITIONAL leg (beyond degree+1) left unconnected, which becomes the + contraction's output index. + + All legs of a CopyNode share one hyperedge symbol, so leaving the extra leg + dangling means the spin variable s_free appears in the einsum output -- for + each of its two values, cotengra maximizes over the other spins. cfg=0 -> + s=+1, cfg=1 -> s=-1 (matching ``build_ising_tn``'s ``Tv=[h,-h]`` convention). + + Returns ``(nodes, free_edge, j_vals, h)`` where ``free_edge`` is the dangling + CopyNode leg to pass as ``output_edge_order=[free_edge]``. + """ + be = tc.backend + degree = dict.fromkeys(spins, 0) + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + + nodes = [] + copy_nodes = {} + for i in spins: + rank = degree[i] + 1 # one field leg + one per incident edge + if i == free_spin: + rank += 1 # extra leg -> dangling output index for the free spin + cn = tn.CopyNode(rank, 2) + copy_nodes[i] = cn + nodes.append(cn) + # field tensor on leg 0 (same as build_ising_tn) + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + + # bond tensors, each consuming one copy leg per endpoint + leg_idx = dict.fromkeys(spins, 1) # leg 0 reserved for field + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) + tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) + leg_idx[i] += 1 + leg_idx[j] += 1 + nodes.append(ten) + + # the free leg is the LAST leg of the free spin's CopyNode + # (legs 0..degree are field+bonds; leg degree+1 is the dangling output). + free_edge = copy_nodes[free_spin][degree[free_spin] + 1] + return nodes, free_edge, j_vals, h + + +def _ising_config_energy(cfg, edges, j_vals, spins, h): + """Ising energy E = -sum J s_i s_j - sum h s_i for a config dict (spin idx -> 0/1).""" + energy = 0.0 + for (i, j), jij in zip(edges, j_vals): + energy -= jij * (1 - 2 * cfg[i]) * (1 - 2 * cfg[j]) + for i, hi in zip(spins, h): + energy -= hi * (1 - 2 * cfg[i]) + return energy + + +def _min_energy_and_degeneracy(energies, eps=1e-9): + """Return ``(min energy, degeneracy)`` -- degeneracy counts configs within eps of min.""" + best = None + deg = 0 + for e in energies: + if best is None or e < best - eps: + best, deg = e, 1 + elif abs(e - best) < eps: + deg += 1 + return best, deg + + +def brute_nonscalar_counting( + nodes, + free_edge, + j_vals, + h, + spins=_RING_SPINS, + edges=_RING_EDGES, + free_spin=_FREE_SPIN, +): + """Brute-force per-output (energy, degeneracy) for the non-scalar ring with + one free spin. Returns ``(expected_e, expected_n)`` -- two arrays of shape + ``(2,)`` indexed by the free spin's value (0 -> s=+1, 1 -> s=-1). + + For each value v of the free spin, enumerate all configs of the OTHER spins, + compute the Ising energy E = -sum J s_i s_j - sum h s_i with s_free fixed, + and record min E + its degeneracy. The contraction returns max-plus(-E) per + output (== -min E), so we return ``-min E`` to match. + + ``nodes`` and ``free_edge`` are accepted only for signature compatibility + with the call site; the brute force is parameterised by the (hardcoded) + ring topology that ``build_ring_with_free_spin`` emits. + """ + del nodes, free_edge # signature-only; topology is fixed by the defaults above + others = [s for s in spins if s != free_spin] + + expected_e = np.zeros(2, dtype=np.float64) + expected_n = np.zeros(2, dtype=np.float64) + for v_free in (0, 1): + fixed = {free_spin: v_free} + energies = [ + _ising_config_energy( + {**fixed, **dict(zip(others, cfg_rest))}, edges, j_vals, spins, h + ) + for cfg_rest in itertools.product([0, 1], repeat=len(others)) + ] + best_e, deg = _min_energy_and_degeneracy(energies) + # max-plus convention: contraction returns max(-E) = -min(E) + expected_e[v_free] = -best_e + expected_n[v_free] = deg + return expected_e, expected_n + + +# --- Task 10 follow-up: non-scalar (two free spins) builder + brute force --- +# +# Extends the one-free-spin canary to a 2D output so the aux-reorder permutation +# in ``cons._algebraic_base_contraction`` (``perm = [dangling_edges.index(e) for e +# in order]``) is exercised with a NON-IDENTITY permutation. With one free edge +# the perm is always ``(0,)`` (identity); with two free edges + a reversed +# ``output_edge_order`` it becomes ``(1, 0)`` (swap), so an elided or +# wrong-direction transpose on the aux tensor is no longer silently a no-op. +# +# Same ring topology / parameters as the one-free-spin case (``_RING_*`` constants +# above), but TWO spins (0 and 1) get an extra dangling CopyNode leg. With +# ``h = [2, 0, 2, -1]`` the per-(a,b) brute-force result is ASYMMETRIC under +# transpose (E = [[5,7],[5,-1]], N = [[1,2],[1,1]]), which is load-bearing: it +# makes ``assert_allclose(N, expected_n_ab.T)`` fail if aux is left in the +# dangling order rather than the ``output_edge_order`` order. +_FREE_SPIN_A = 0 +_FREE_SPIN_B = 1 + + +def build_ring_with_two_free_spins( + spins=_RING_SPINS, + edges=_RING_EDGES, + j_vals=_RING_J, + h=_RING_H, + free_spin_a=_FREE_SPIN_A, + free_spin_b=_FREE_SPIN_B, +): + """Build a small ring Ising with TWO spins' CopyNodes given an extra dangling + leg each (two free output indices). Same ``Tv``/``Te``/``tn.CopyNode`` + construction as ``build_ising_tn`` / ``build_ring_with_free_spin`` -- the only + difference is that BOTH ``free_spin_a`` and ``free_spin_b`` have a CopyNode of + rank ``degree + 2`` (one field leg + one per incident edge + one dangling + output leg). + + The two dangling legs become the contraction's two-axis output. Under + ``cons.sorted_edges`` the dangling order comes out as ``[edge_a, edge_b]`` + (spin 0 before spin 1), so contracting with + ``output_edge_order=[edge_b, edge_a]`` exercises a NON-trivial aux + permutation ``(1, 0)``. + + Returns ``(nodes, edge_a, edge_b, j_vals, h)`` where ``edge_a`` is the + dangling leg of ``free_spin_a`` (spin 0) and ``edge_b`` that of + ``free_spin_b`` (spin 1). Pass them to ``output_edge_order`` in whichever + order the test wants to validate. + """ + be = tc.backend + free_spins = {free_spin_a, free_spin_b} + degree = dict.fromkeys(spins, 0) + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + + nodes = [] + copy_nodes = {} + for i in spins: + rank = degree[i] + 1 # one field leg + one per incident edge + if i in free_spins: + rank += 1 # extra leg -> dangling output index for this free spin + cn = tn.CopyNode(rank, 2) + copy_nodes[i] = cn + nodes.append(cn) + # field tensor on leg 0 (same as build_ising_tn) + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + + # bond tensors, each consuming one copy leg per endpoint + leg_idx = dict.fromkeys(spins, 1) # leg 0 reserved for field + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) + tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) + leg_idx[i] += 1 + leg_idx[j] += 1 + nodes.append(ten) + + # each free leg is the LAST leg of its CopyNode + # (legs 0..degree are field+bonds; leg degree+1 is the dangling output). + edge_a = copy_nodes[free_spin_a][degree[free_spin_a] + 1] + edge_b = copy_nodes[free_spin_b][degree[free_spin_b] + 1] + return nodes, edge_a, edge_b, j_vals, h + + +def brute_nonscalar_counting_2d( + nodes, + edge_a, + edge_b, + j_vals, + h, + spins=_RING_SPINS, + edges=_RING_EDGES, + free_spin_a=_FREE_SPIN_A, + free_spin_b=_FREE_SPIN_B, +): + """Brute-force per-output ``(energy, degeneracy)`` for the non-scalar ring + with TWO free spins. Returns ``(expected_e, expected_n)`` -- two ``(2, 2)`` + arrays indexed by ``[v_a, v_b]`` where ``v_a`` is the value of + ``free_spin_a`` (cfg=0 -> s=+1, cfg=1 -> s=-1) and ``v_b`` that of + ``free_spin_b``. + + For each ``(v_a, v_b)`` pair, enumerate the other spins, compute + ``E = -sum J s_i s_j - sum h s_i`` with both free spins fixed, and record + ``min E`` + its degeneracy. Returns ``-min E`` to match the max-plus sign + convention of ``brute_nonscalar_counting``. + + ``nodes`` / ``edge_a`` / ``edge_b`` are signature-only (the topology is fixed + by the defaults above); they are accepted so the call site reads symmetrically + with ``build_ring_with_two_free_spins``'s return signature. + + For the default parameters the brute-force result is + ``E = [[5, 7], [5, -1]]`` and ``N = [[1, 2], [1, 1]]``, both asymmetric under + transpose -- a property the test relies on so a missed aux transpose cannot + pass ``assert_allclose(N, expected_n_ab.T)``. + """ + del nodes, edge_a, edge_b # signature-only; topology is fixed by the defaults above + others = [s for s in spins if s not in (free_spin_a, free_spin_b)] + + expected_e = np.zeros((2, 2), dtype=np.float64) + expected_n = np.zeros((2, 2), dtype=np.float64) + for v_a in (0, 1): + for v_b in (0, 1): + fixed = {free_spin_a: v_a, free_spin_b: v_b} + energies = [ + _ising_config_energy( + {**fixed, **dict(zip(others, cfg_rest))}, edges, j_vals, spins, h + ) + for cfg_rest in itertools.product([0, 1], repeat=len(others)) + ] + best_e, deg = _min_energy_and_degeneracy(energies) + # max-plus convention: contraction returns max(-E) = -min(E) + expected_e[v_a, v_b] = -best_e + expected_n[v_a, v_b] = deg + return expected_e, expected_n diff --git a/tests/test_bcomplex32_algebra.py b/tests/test_bcomplex32_algebra.py new file mode 100644 index 00000000..3c584712 --- /dev/null +++ b/tests/test_bcomplex32_algebra.py @@ -0,0 +1,90 @@ +import numpy as np +from applications.bcomplex32_algebra import ( + _pair_tensordot, + _pair_einsum, + _complex_to_pair, + _pair_to_complex, +) +import tensorcircuit.backends.numpy_backend as nb + +be = nb.NumpyBackend() + + +def test_complex_to_pair_roundtrip_within_bf16_quant(): + t = np.array( + [[1.0 + 2.0j, 3.0 - 1.0j], [0.25 + 0.5j, -2.0 + 7.0j]], dtype=np.complex64 + ) + pair = _complex_to_pair(be, t) + back = _pair_to_complex(be, pair) + np.testing.assert_allclose(np.asarray(back), t, rtol=2e-2) + + +def test_pair_tensordot_matches_complex_tensordot(): + rng = np.random.default_rng(0) + a = rng.standard_normal((2, 3)).astype(np.complex64) + b = rng.standard_normal((3, 4)).astype(np.complex64) + pa = _complex_to_pair(be, a) + pb = _complex_to_pair(be, b) + out = _pair_tensordot(be, pa, pb, axes=([1], [0])) + ref = np.tensordot(a, b, axes=([1], [0])) + np.testing.assert_allclose( + np.asarray(_pair_to_complex(be, out)), ref, rtol=2e-2, atol=5e-3 + ) + + +def test_pair_einsum_single_operand(): + rng = np.random.default_rng(0) + a = rng.standard_normal((2, 2)).astype(np.complex64) + pa = _complex_to_pair(be, a) + out = _pair_einsum( + be, "ab->a", pa + ) # reduce over b (sum in complex -> here identity semantics) + ref = np.einsum("ab->a", a) + np.testing.assert_allclose( + np.asarray(_pair_to_complex(be, out)), ref, rtol=2e-2, atol=5e-3 + ) + + +# --- Task 12: end-to-end through real tc.Circuit + wall-avoidance canary --- + +import tensorcircuit as tc +from applications.bcomplex32_algebra import bcomplex32 + + +def test_bf16_end_to_end_matches_complex64(): + import numpy as np + + def build(): + c = tc.Circuit(3) + c.H(0) + c.cnot(0, 1) + c.cnot(1, 2) + return np.asarray(c.state()) + + ref = build() # default complex64 + with bcomplex32(): + got = build() # bf16 pair path + np.testing.assert_allclose(got, ref, rtol=2e-2) + + +def test_bf16_wall_avoidance_canary(): + import numpy as np + + with bcomplex32(): + c = tc.Circuit(4) + for i in range(4): + c.H(i) + for i in range(3): + 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 + from tensorcircuit.contraction_algebra import StandardAlgebra + + assert isinstance( + cons.get_contraction_algebra(), StandardAlgebra + ) # CM restored algebra, no leak + c2 = tc.Circuit(2) + c2.H(0) + c2.cnot(0, 1) # subsequent native contraction + assert np.asarray(c2.state()).shape == (4,) # algebra restored, no leak diff --git a/tests/test_contraction_algebra.py b/tests/test_contraction_algebra.py new file mode 100644 index 00000000..c325a91c --- /dev/null +++ b/tests/test_contraction_algebra.py @@ -0,0 +1,409 @@ +import numpy as np +import tensorcircuit as tc +import tensorcircuit.cons as cons +from tensorcircuit.contraction_algebra.base import ContractionAlgebra, StandardAlgebra + + +def test_standard_tensordot_matches_backend(): + be = tc.backend + a = be.cast( + be.convert_to_tensor(np.arange(12, dtype=np.float64).reshape(3, 4)), "float64" + ) + b = be.cast( + be.convert_to_tensor(np.arange(20, dtype=np.float64).reshape(4, 5)), "float64" + ) + alg = StandardAlgebra() + got = np.array(alg.tensordot(be, a, b, axes=1)) + ref = np.array(be.tensordot(a, b, axes=1)) + assert got.shape == (3, 5) + np.testing.assert_allclose(got, ref) + + +def test_standard_einsum_matches_backend(): + be = tc.backend + a = be.cast( + be.convert_to_tensor(np.arange(6, dtype=np.float64).reshape(2, 3)), "float64" + ) + b = be.cast( + be.convert_to_tensor(np.arange(12, dtype=np.float64).reshape(3, 4)), "float64" + ) + alg = StandardAlgebra() + got = np.array(alg.einsum(be, "ab,bc->ac", a, b)) + ref = np.array(be.einsum("ab,bc->ac", a, b)) + np.testing.assert_allclose(got, ref) + + +def test_standard_name(): + assert StandardAlgebra().name == "standard" + + +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 + # in-source via cons.set_contraction_algebra / set_contractor(algebra=...). + for name in [ + "ContractionAlgebra", + "StandardAlgebra", + "Representation", + "IdentityRepresentation", + ]: + assert hasattr(tca, name), name + # The old monkey-patch API names are intentionally gone: + for gone in [ + "activate", + "deactivate", + "standard", + "runtime_contraction_algebra", + "set_contraction_algebra", + "get_contraction_algebra", + "injection", + ]: + assert not hasattr(tca, gone), gone + + +def test_algebra_hooks_default_noop(): + alg = StandardAlgebra() + assert alg.on_contraction_start(["dummy_nodes"]) is None + assert alg.on_contractor_ready(["dummy_tree"]) is None + + +from tensorcircuit.contraction_algebra import ( + ContractionAlgebra, + StandardAlgebra, + Representation, + IdentityRepresentation, +) + +# --- Task 5: non-standard path (encode -> kernels via implementation= -> decode, +# hooks fire, primary.ndim == len(output_set), aux stashed) --- + + +def test_nonstandard_path_encodes_kernel_decodes_in_order(): + import numpy as np + import tensornetwork as tn + import opt_einsum + + log = [] + + class LogRep(Representation): + name = "log" + + def encode(self, be, tensors): + log.append("encode") + return tensors + + def decode(self, be, tensor): + log.append("decode") + return tensor, {} + + class LogAlg(ContractionAlgebra): + name = "log" + representation = LogRep() + + def tensordot(self, be, a, b, axes): + log.append("td") + return be.tensordot(a, b, axes) + + def einsum(self, be, eq, *ops): + log.append("ein") + return be.einsum(eq, *ops) + + def on_contraction_start(self, nodes): + log.append("start") + + def on_contractor_ready(self, tree): + log.append("ready") + + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal((2, 3)).astype(np.complex64)) + b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) + tn.connect(a[1], b[0]) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(LogAlg()) + try: + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], b[1]], + ) + assert log[0] == "start" + assert "encode" in log and "ready" in log and "decode" in log + assert log.index("encode") < log.index("decode") + # hooks fire around encode/ready in the prescribed order: + assert log.index("start") < log.index("encode") + assert log.index("ready") < log.index("decode") + finally: + cons.set_contraction_algebra(prev) + + +def test_decode_wrong_rank_raises_loudly(): + import numpy as np + import tensornetwork as tn + import opt_einsum + import pytest + + class BadRep(Representation): + name = "bad" + + def encode(self, be, tensors): + return tensors + + def decode(self, be, tensor): + # Add a trailing storage axis that the representation forgot to strip. + return ( + be.reshape(tensor, tuple(be.shape_tuple(tensor)) + (1,)), + {}, + ) + + class BadAlg(ContractionAlgebra): + name = "bad" + representation = BadRep() + + def tensordot(self, be, a, b, axes): + return be.tensordot(a, b, axes) + + def einsum(self, be, eq, *ops): + return be.einsum(eq, *ops) + + # 2-node fully contracted scalar: output_set="", len 0. BadAlg.decode returns + # rank 1, so primary.ndim != len(output_set) trips the assert. + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal(2).astype(np.complex64)) + b = tn.Node(rng.standard_normal(2).astype(np.complex64)) + tn.connect(a[0], b[0]) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(BadAlg()) + try: + with pytest.raises(AssertionError): + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[], + ) + finally: + cons.set_contraction_algebra(prev) + + +def test_aux_outputs_stashed_and_reordered(): + # The full aux-reorder test is Task 8 (counting). Here we only confirm the + # side-channel store API exists and round-trips a value. + cons._stash_aux_outputs({"count": 5}) + assert cons._aux_outputs()["count"] == 5 + + +def test_representation_identity_roundtrip(): + rep = IdentityRepresentation() + import numpy as np + + t = [np.ones((2, 3))] + assert rep.encode(None, t) is t # identity returns same + primary, aux = rep.decode(None, t[0]) + assert primary is t[0] and aux == {} + + +def test_standard_algebra_carries_identity_representation(): + alg = StandardAlgebra() + assert isinstance(alg.representation, IdentityRepresentation) + assert alg.name == "standard" + + +def test_contraction_algebra_representation_default_is_identity(): + # a minimal concrete algebra that does NOT override representation + class BareAlgebra(ContractionAlgebra): + name = "bare" + + def tensordot(self, be, a, b, axes): + return be.tensordot(a, b, axes) + + def einsum(self, be, eq, *ops): + return be.einsum(eq, *ops) + + assert isinstance(BareAlgebra().representation, IdentityRepresentation) + + +# --- Task 2: cons.py algebra state + set_contractor(algebra=) --- + + +class _NS(ContractionAlgebra): # non-standard stub for constraint tests + name = "ns" + representation = None # noqa — not used in these tests + + def tensordot(self, be, a, b, axes): + return be.tensordot(a, b, axes) + + def einsum(self, be, eq, *ops): + return be.einsum(eq, *ops) + + +def test_set_contractor_accepts_algebra_kwarg(): + prev = cons.get_contraction_algebra() + try: + cons.set_contractor("greedy", algebra=_NS()) + assert isinstance(cons.get_contraction_algebra(), _NS) + finally: + cons.set_contraction_algebra(prev) + + +def test_set_contractor_rejects_preprocessing_with_nonstandard_algebra(): + import pytest + + ns_algebra = _NS() + with pytest.raises(ValueError): + cons.set_contractor("greedy", algebra=ns_algebra, preprocessing=True) + + +def test_set_contractor_rejects_strip_exponent_with_nonstandard_algebra(): + import pytest + + ns_algebra = _NS() + with pytest.raises(ValueError): + cons.set_contractor("greedy", algebra=ns_algebra, strip_exponent=True) + + +def test_runtime_contractor_restores_algebra(): + prev = cons.get_contraction_algebra() + with cons.runtime_contractor("greedy", algebra=_NS()): + assert isinstance(cons.get_contraction_algebra(), _NS) + assert cons.get_contraction_algebra() is prev + + +def test_set_function_contractor_restores_algebra(): + prev = cons.get_contraction_algebra() + + @cons.set_function_contractor("greedy", algebra=_NS()) + def f(): + return cons.get_contraction_algebra() + + inside = f() + assert isinstance(inside, _NS) # algebra active during the call + assert cons.get_contraction_algebra() is prev # restored after + + +# --- Task 3: cons.py guards (_merge_single_gates skip + _base routing) --- + + +def test_merge_single_gates_skipped_under_nonstandard_algebra(monkeypatch): + import tensornetwork as tn + + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(_NS()) + try: + + def boom(*a, **k): + raise AssertionError( + "merge body ran — should be skipped under non-standard algebra" + ) + + monkeypatch.setattr(tn, "contract_parallel", boom) + out = cons._merge_single_gates( + ["fake_node"], 7 + ) # guard fires before any node access + assert out == (["fake_node"], 7) + finally: + cons.set_contraction_algebra(prev) + + +def test_base_routes_to_algebraic_under_nonstandard(monkeypatch): + import numpy as np + import tensornetwork as tn + import opt_einsum + + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal((2, 3)).astype(np.complex64)) + b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) + tn.connect(a[1], b[0]) + routed = {} + + def fake_alg(nodes, algorithm, *args, **kw): + routed["called"] = True + return "FINAL" + + monkeypatch.setattr(cons, "_algebraic_base_contraction", fake_alg) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(_NS()) + try: + cons._base( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], b[1]], + ) + assert routed.get("called") is True + finally: + cons.set_contraction_algebra(prev) + + +# --- Task 4: lock the StandardAlgebra backward-compat baseline --- +# A real 2-node contraction under the default StandardAlgebra must produce +# results bit-identical to a direct np.tensordot reference. This is the +# safety net before Task 5 adds the non-standard encode/decode path to +# _algebraic_base_contraction. If this test passes on today's code (which +# ignores the algebra), that is exactly the point: it locks the baseline. + + +def test_standard_algebra_contraction_matches_native(): + import numpy as np + import tensornetwork as tn + + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal((2, 3)).astype(np.complex64)) + b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) + tn.connect(a[1], b[0]) + prev = cons.get_contraction_algebra() + assert isinstance(prev, StandardAlgebra) + cons.set_contraction_algebra(StandardAlgebra()) + try: + import opt_einsum + + n = cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], b[1]], + ) + ref = np.tensordot(a.tensor, b.tensor, axes=([1], [0])) + np.testing.assert_allclose(n.tensor, ref, rtol=1e-6) + finally: + cons.set_contraction_algebra(prev) + + +def test_aux_stash_handles_ignore_edge_order_with_none_order(monkeypatch): + # A non-standard algebra returning aux under ignore_edge_order=True + output_edge_order=None + # must NOT crash (counting-scalar path). Uses the LogAlg/LogRep from test_nonstandard_path_*. + import numpy as np + import tensornetwork as tn + import opt_einsum + + class AuxRep(Representation): + def encode(self, be, tensors): + return tensors + + def decode(self, be, tensor): + return tensor, {"count": np.ones_like(tensor)} # non-empty aux + + class AuxAlg(ContractionAlgebra): + name = "aux" + representation = AuxRep() + + def tensordot(self, be, a, b, axes): + return be.tensordot(a, b, axes) + + def einsum(self, be, eq, *ops): + return be.einsum(eq, *ops) + + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal(2).astype(np.complex64)) + b = tn.Node(rng.standard_normal(2).astype(np.complex64)) + tn.connect(a[0], b[0]) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(AuxAlg()) + try: + # scalar contraction, ignore_edge_order=True, output_edge_order=None — must not raise + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=None, + ignore_edge_order=True, + ) + assert "count" in cons._aux_outputs() # aux stashed without crashing + finally: + cons.set_contraction_algebra(prev) diff --git a/tests/test_tropical_algebra.py b/tests/test_tropical_algebra.py new file mode 100644 index 00000000..05ab7da0 --- /dev/null +++ b/tests/test_tropical_algebra.py @@ -0,0 +1,272 @@ +import numpy as np +import tensornetwork as tn +import tensorcircuit as tc +import tensorcircuit.cons as cons +from applications.tropical_algebra import _tropical_tensordot + + +def _be(): + return tc.backend + + +def _ref_tropical_tensordot(anp, bnp, axes): + if isinstance(axes, int): + a_axes = list(range(anp.ndim - axes, anp.ndim)) + b_axes = list(range(0, axes)) + else: + a_axes, b_axes = list(axes[0]), list(axes[1]) + a_free = [i for i in range(anp.ndim) if i not in a_axes] + b_free = [i for i in range(bnp.ndim) if i not in b_axes] + a_t = np.transpose(anp, a_free + a_axes) + b_t = np.transpose(bnp, b_axes + b_free) + a_fs = [anp.shape[i] for i in a_free] + b_fs = [bnp.shape[i] for i in b_free] + a2 = a_t.reshape(-1, np.prod([anp.shape[i] for i in a_axes], dtype=int) or 1) + b2 = b_t.reshape(np.prod([bnp.shape[i] for i in b_axes], dtype=int) or 1, -1) + A, b_n = a2.shape[0], b2.shape[1] + res = np.full((A, b_n), -np.inf) + for i in range(A): + for j in range(b_n): + res[i, j] = np.max(a2[i, :] + b2[:, j]) + return res.reshape(tuple(a_fs) + tuple(b_fs)) + + +def test_tropical_tensordot_matrix(): + be = _be() + rng = np.random.default_rng(0) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_tensordot(be, a, b, axes=1)) + np.testing.assert_allclose(got, _ref_tropical_tensordot(anp, bnp, 1)) + + +def test_tropical_tensordot_multi_axis(): + be = _be() + rng = np.random.default_rng(1) + anp, bnp = rng.normal(size=(2, 3, 4)), rng.normal(size=(3, 4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_tensordot(be, a, b, axes=([1, 2], [0, 1]))) + np.testing.assert_allclose(got, _ref_tropical_tensordot(anp, bnp, ([1, 2], [0, 1]))) + assert got.shape == (2, 5) + + +from applications.tropical_algebra import _tropical_einsum, MaxPlusAlgebra, tropical + + +def _ref_tropical_einsum(eq, a, b): + import itertools + + lhs, rhs = eq.split("->") + ia, ib = lhs.split(",") + sizes = {} + for s, t in zip([ia, ib], [a, b]): + for c, dim in zip(s, t.shape): + sizes[c] = dim + out = np.full([sizes[c] for c in rhs], -np.inf) + allc = list(dict.fromkeys(list(ia) + list(ib))) + for combo in itertools.product(*[range(sizes[c]) for c in allc]): + env = dict(zip(allc, combo)) + ia_idx = tuple(env[c] for c in ia) + ib_idx = tuple(env[c] for c in ib) + val = a[ia_idx] + b[ib_idx] + oidx = tuple(env[c] for c in rhs) + if val > out[oidx]: + out[oidx] = val + return out + + +def test_tropical_einsum_pair(): + be = _be() + rng = np.random.default_rng(2) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "ab,bc->ac", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("ab,bc->ac", anp, bnp)) + + +def test_tropical_einsum_hyperedge_shared_index(): + # 'a' shared (batch, kept in output) — this is the copy-node pair shape + be = _be() + rng = np.random.default_rng(3) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "ab,ac->abc", a, b)) # note: output order a,b,c + ref = _ref_tropical_einsum("ab,ac->abc", anp, bnp) + np.testing.assert_allclose(got, ref) + + +def test_maxplus_algebra_uses_tropical_ops(): + be = _be() + rng = np.random.default_rng(4) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + alg = MaxPlusAlgebra() + assert alg.name == "maxplus" + np.testing.assert_allclose( + np.array(alg.tensordot(be, a, b, 1)), + _ref_tropical_tensordot(anp, bnp, 1), + ) + np.testing.assert_allclose( + np.array(alg.einsum(be, "ab,bc->ac", a, b)), + _ref_tropical_einsum("ab,bc->ac", anp, bnp), + ) + + +def test_tropical_einsum_single_tensor_transpose_ok(): + # single tensor with NO repeated indices -> standard transpose is correct + be = _be() + rng = np.random.default_rng(5) + anp = rng.normal(size=(3, 4)) + a = be.cast(be.convert_to_tensor(anp), "float64") + got = np.array(_tropical_einsum(be, "ab->ba", a)) + np.testing.assert_allclose(got, np.array(be.einsum("ab->ba", a))) + + +def test_tropical_einsum_single_tensor_trace(): + # "ii->" : tropical trace = max of the diagonal + be = _be() + rng = np.random.default_rng(6) + square_a = rng.normal(size=(3, 3)) + a = be.cast(be.convert_to_tensor(square_a), "float64") + got = float(np.array(_tropical_einsum(be, "ii->", a))) + assert np.isclose(got, np.max(np.diag(square_a))) + + +def test_tropical_einsum_single_tensor_diagonal(): + # "ii->i" : tropical diagonal gather (repeated index kept, no reduction) + be = _be() + rng = np.random.default_rng(7) + square_a = rng.normal(size=(4, 4)) + a = be.cast(be.convert_to_tensor(square_a), "float64") + got = np.array(_tropical_einsum(be, "ii->i", a)) + np.testing.assert_allclose(got, np.diag(square_a)) + + +def test_tropical_einsum_single_tensor_reduce(): + # "ab->a" : single-tensor axis reduction (no repeat) must be tropical max, not sum + be = _be() + rng = np.random.default_rng(8) + anp = rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + got = np.array(_tropical_einsum(be, "ab->a", a)) + np.testing.assert_allclose(got, np.max(anp, axis=1)) + + +def test_tropical_einsum_single_tensor_partial_trace(): + # "iij->j" : repeated index reduced (diagonal over i, then max) + be = _be() + rng = np.random.default_rng(9) + anp = rng.normal(size=(4, 4, 3)) # axes: i, i, j + a = be.cast(be.convert_to_tensor(anp), "float64") + got = np.array(_tropical_einsum(be, "iij->j", a)) + ref = np.max(np.array([anp[i, i, :] for i in range(4)]), axis=0) + np.testing.assert_allclose(got, ref) + + +def test_tropical_einsum_intra_operand_repeat_first(): + # "iij,jk->ik": first operand has a repeated index i (diagonal gather), then + # ordinary pairwise tropical contraction over j. Spec §4.5 step 1. + be = _be() + rng = np.random.default_rng(20) + anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "iij,jk->ik", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jk->ik", anp, bnp)) + + +def test_tropical_einsum_intra_operand_repeat_contracted(): + # "iij,jk->k": the repeated index i is neither in the output nor shared with + # b -> diagonal gather of i, then tropical max over both i and j. + be = _be() + rng = np.random.default_rng(21) + anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "iij,jk->k", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jk->k", anp, bnp)) + + +def test_tropical_einsum_both_operands_repeat(): + # "iij,jj->ij": both operands have intra-operand repeats; no contraction axis. + be = _be() + rng = np.random.default_rng(22) + anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 4)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "iij,jj->ij", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jj->ij", anp, bnp)) + + +def test_tropical_context_changes_result(): + # max-plus "ab,b->a" = max_b (A[a,b] + B[b]) differs from standard sum-product + be = _be() + anp = np.array([[1.0, 5.0], [3.0, 2.0]]) + bnp = np.array([10.0, 0.0]) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + na, nb = tn.Node(a), tn.Node(b) + tn.connect(na[1], nb[0]) + expected = np.array([max(1 + 10, 5 + 0), max(3 + 10, 2 + 0)]) # [11, 13] + with tropical(): + got = np.array(cons.contractor([na, nb], output_edge_order=[na[0]]).tensor) + np.testing.assert_allclose(got, expected) + + +def test_preprocessing_does_not_corrupt_tropical(): + # Regression: default cons.contractor (greedy + preprocessing=True) calls + # _merge_single_gates for hyperedge-free >=5-node networks, merging single-gate + # nodes via STANDARD tn.contract_parallel BEFORE the algebraic path runs. For a + # tropical TN of regular tn.Nodes (no CopyNode -> no hyperedge -> preprocessing + # not skipped), that merge would corrupt the max-plus result with sum-product. + # The in-source _merge_single_gates guard skips the merge under a non-standard algebra. + import itertools + + be = _be() + rng = np.random.default_rng(0) + n0 = rng.normal(size=(2,)) + n1 = rng.normal(size=(2, 2)) + n2 = rng.normal(size=(2, 2)) + n3 = rng.normal(size=(2, 2)) + n4 = rng.normal(size=(2,)) + na = tn.Node(be.cast(be.convert_to_tensor(n0), "float64")) + nb = tn.Node(be.cast(be.convert_to_tensor(n1), "float64")) + nc = tn.Node(be.cast(be.convert_to_tensor(n2), "float64")) + nd = tn.Node(be.cast(be.convert_to_tensor(n3), "float64")) + ne = tn.Node(be.cast(be.convert_to_tensor(n4), "float64")) + tn.connect(na[0], nb[0]) + tn.connect(nb[1], nc[0]) + tn.connect(nc[1], nd[0]) + tn.connect(nd[1], ne[0]) + nodes = [na, nb, nc, nd, ne] + assert len(nodes) >= 5 + + ref = -np.inf + for s0, s1, s2, s3 in itertools.product(range(2), repeat=4): + val = n0[s0] + n1[s0, s1] + n2[s1, s2] + n3[s2, s3] + n4[s3] + ref = max(ref, val) + + with tropical(): + got = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + + np.testing.assert_allclose(got, ref, atol=1e-6) + + +def test_tropical_public_api_surface(): + import applications.tropical_algebra as ex + + for name in [ + "MaxPlusAlgebra", + "MaxPlusTrackingAlgebra", + "CountingTropicalAlgebra", + "tropical", + "counting_tropical", + "recover_configuration", + "split_energy_count", + ]: + assert hasattr(ex, name), name diff --git a/tests/test_tropical_config.py b/tests/test_tropical_config.py new file mode 100644 index 00000000..13b64b79 --- /dev/null +++ b/tests/test_tropical_config.py @@ -0,0 +1,246 @@ +"""Task A1 (SPIKE) validation: argmax-tracking + tree-backtracking recovers an +optimal configuration of a tiny tropical (max, +) Ising contraction. + +The assertion is energy-based: the configuration returned by +``recover_configuration()`` must achieve the brute-force optimum energy. It need +not equal a *specific* brute-force argmin when the ground state is degenerate, +but its energy must be optimal. +""" + +import itertools + +import numpy as np +import pytest +import tensornetwork as tn + +import tensorcircuit as tc +import tensorcircuit.cons as cons +from applications.tropical_algebra import ( + tropical, + recover_configuration, + get_recorded_topology, + _td_backtrack_step, +) +from tests._tropical_test_utils import build_ising_tn + + +def _brute_cfg(n, edges, j_vals, h): + """Brute-force min energy over spin configs; returns (best_e, set_of_optimal_cfgs). + + cfg in {0,1}^n; spin s = 1 - 2*cfg (cfg=0 -> s=+1). E = -sum J s_i s_j - sum h s_i. + """ + best_e = None + best_cfgs = set() + for cfg in itertools.product([0, 1], repeat=n): + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + e -= jij * (1 - 2 * cfg[i]) * (1 - 2 * cfg[j]) + for i, hi in zip(range(n), h): + e -= hi * (1 - 2 * cfg[i]) + if best_e is None or e < best_e - 1e-9: + best_e = e + best_cfgs = {cfg} + elif abs(e - best_e) < 1e-9: + best_cfgs.add(cfg) + return best_e, best_cfgs + + +def _energy_of(assignment, input_sets, raw_tensors): + """Max-plus total for a symbol->value assignment = -E for that config.""" + total = 0.0 + for term, t in zip(input_sets, raw_tensors): + idx = tuple(int(assignment[c]) for c in term) + total += float(np.asarray(t)[idx]) + return total + + +def _contract_tropical_track(nodes): + """Contract under tracking tropical; return (value, config).""" + with tropical(track=True): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + cfg = recover_configuration() + return val, cfg + + +def test_config_recovery_tiny_ising_unique(): + # 2-spin Ising: spin0 --J-- spin1, fields h0, h1 -> unique ground state. + spins = [0, 1] + edges = [(0, 1)] + J = [0.7] + h = [0.3, -0.2] + best_e, best_cfgs = _brute_cfg(len(spins), edges, J, h) + + nodes = build_ising_tn(spins, edges, J, h) + val, cfg = _contract_tropical_track(nodes) + + # contraction returns max_cfg(-E) = -E_ground + np.testing.assert_allclose(val, -best_e, atol=1e-6) + + _tree, input_sets, raw_tensors = get_recorded_topology() + assert input_sets is not None and raw_tensors is not None + # every index label got a value + all_labels = set().union(*[set(t) for t in input_sets]) + assert set(cfg.keys()) == all_labels + + recovered_total = _energy_of(cfg, input_sets, raw_tensors) + # the recovered config's max-plus total must equal the contraction value + # and hence -best_e (optimal energy). + np.testing.assert_allclose(recovered_total, val, atol=1e-6) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + # unique ground state: the recovered (symbol) config must map to one of the + # brute-force optimal cfgs. Map symbols -> per-spin cfg via the field tensors: + # the field tensor for spin i has a single symbol whose value is cfg[i]. + # Recover that mapping from the topology (field tensors are the rank-1 terms). + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: # field tensor -> one spin + spin_cfg[len(spin_cfg)] = int(cfg[term[0]]) + assert tuple(spin_cfg[i] for i in range(len(spins))) in best_cfgs + + +def test_config_recovery_tiny_ising_degenerate(): + # Degenerate ground state (Z2 spin-flip symmetry): ferromagnetic bond, no field. + # Any recovered config must still be optimal (energy-based assertion). + spins = [0, 1] + edges = [(0, 1)] + J = [1.0] + h = [0.0, 0.0] + best_e, best_cfgs = _brute_cfg(len(spins), edges, J, h) + assert len(best_cfgs) == 2 # all-up and all-down + + nodes = build_ising_tn(spins, edges, J, h) + val, cfg = _contract_tropical_track(nodes) + + np.testing.assert_allclose(val, -best_e, atol=1e-6) + _tree, input_sets, raw_tensors = get_recorded_topology() + recovered_total = _energy_of(cfg, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg[term[0]]) + assert tuple(spin_cfg[i] for i in range(len(spins))) in best_cfgs + + +def test_config_recovery_three_ring(): + # Slightly bigger: 3-spin open chain -> exercises a 3rd pairwise step. + spins = [0, 1, 2] + edges = [(0, 1), (1, 2)] + J = [0.5, -0.8] + h = [0.1, 0.2, -0.3] + best_e, _best_cfgs = _brute_cfg(len(spins), edges, J, h) + + nodes = build_ising_tn(spins, edges, J, h) + val, cfg = _contract_tropical_track(nodes) + np.testing.assert_allclose(val, -best_e, atol=1e-6) + + _tree, input_sets, raw_tensors = get_recorded_topology() + recovered_total = _energy_of(cfg, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + +def test_config_recovery_non_scalar_raises(): + """``recover_configuration()`` is scalar-only: a contraction whose + root has dangling/free output indices (non-scalar result) must raise + ``NotImplementedError`` rather than return a (wrong) configuration. + + Builds a max-plus matrix-vector product + ``result[i] = max_k(A[i,k] + B[k])`` (one free index ``i``, one contracted + index ``k``) -- a non-scalar root -- and asserts recovery raises. The + non-scalar backtracking wiring has known ordering bugs (tree.output order + vs result-shape order; output-label values lost in the einsum branch), so it + is gated instead of shipping wrong answers. + """ + be = tc.backend + a_np = np.array( + [[1.0, 4.0, 2.0], [3.0, 1.0, 5.0], [0.0, 2.0, 1.0]] + ) # A[i,k], shape (3,3) + b_np = np.array([0.5, 1.0, -0.5]) # B[k], shape (3,) + + a_node = tn.Node(be.cast(be.convert_to_tensor(a_np), "float64")) + b_node = tn.Node(be.cast(be.convert_to_tensor(b_np), "float64")) + tn.connect(a_node[1], b_node[0]) # contract k; leave i (a_node[0]) dangling + nodes = [a_node, b_node] + + with tropical(track=True): + res = cons.contractor(nodes, output_edge_order=[a_node[0]]) + # sanity: the contraction itself is fine and non-scalar (shape (3,)) + assert np.array(res.tensor).shape == (3,) + with pytest.raises(NotImplementedError): + recover_configuration() + + +class _StubTree: + """Minimal stand-in for a cotengra ``ContractionTree`` exposing only the + three methods that ``_td_backtrack_step`` reads (``get_inds``, + ``get_tensordot_axes``, ``get_tensordot_perm``). Used to feed the backtrack + step a synthetic non-``None`` perm without spinning up a real contraction. + """ + + def __init__(self, inds, axes, perm): + self._inds = inds + self._axes = axes + self._perm = perm + + def get_inds(self, node): + return self._inds[node] + + def get_tensordot_axes(self, node): + return self._axes[node] + + def get_tensordot_perm(self, node): + return self._perm[node] + + +def test_td_backtrack_step_perm_inversion(): + """Synthetic-perm unit test for ``_td_backtrack_step`` (the I1 guard). + + The scalar 5-ring canary contracts with all-``None`` ``get_tensordot_perm`` + values, so the perm!=None inversion branch -- mapping a canonical + (``get_inds(p)``) position to the algebra-natural (``l_free + r_free``) + order via ``td_pos[perm[i]] = p_pos[i]`` -- is NEVER exercised end-to-end. + This feeds the backtrack step a synthetic 3-cycle perm (non-self-inverse, so + a wrong inversion *direction* is caught, unlike a transposition) and asserts + the recovered per-axis positions are exactly the inverse of the perm. + + Guards the inversion logic for future non-scalar support and against cotengra + convention changes. + """ + # Geometry: l = 'ab' (a free, b contracted), r = 'bcd' (b contracted, c,d + # free). l_axes=[1] (b is axis 1 of l), r_axes=[0] (b is axis 0 of r). + # Algebra-natural output order = l_free + r_free = 'acd'. + # Canonical (get_inds(p)) order chosen as 'dac' (a 3-cycle of 'acd'), so the + # perm is non-trivial and non-self-inverse: + # perm[i] = 'acd'.find('dac'[i]) = (2, 0, 1). + l_node, r_node, p_node = "l", "r", "p" + tree = _StubTree( + inds={l_node: "ab", r_node: "bcd", p_node: "dac"}, + axes={p_node: ([1], [0])}, + perm={p_node: (2, 0, 1)}, + ) + # argmax lives in algebra-natural 'acd' order, shape (dim_a, dim_c, dim_d). + # Place flattened contracted index 3 at td_pos (a=1, c=2, d=0); with + # contract_dims=(4,) -> contracted label 'b' value 3. + argmax = np.zeros((2, 3, 2), dtype=np.int64) + argmax[1, 2, 0] = 3 + rec = {"kind": "td", "argmax": argmax, "contract_dims": (4,)} + # Canonical p_pos in 'dac' order: d=0, a=1, c=2. + p_pos = (0, 1, 2) + assignment = {} + l_pos, r_pos = _td_backtrack_step( + tree, p_node, l_node, r_node, rec, p_pos, assignment + ) + + # Inversion applied: canonical 'dac'(0,1,2) -> algebra-natural 'acd'(1,2,0), + # i.e. td_pos = (a=1, c=2, d=0); argmax[1,2,0]=3 -> 'b'=3. + assert assignment["b"] == 3 + # l_pos in 'ab' order: a=1 (free via inversion), b=3 (contracted). + assert l_pos == (1, 3) + # r_pos in 'bcd' order: b=3 (contracted), c=2 (free), d=0 (free). + assert r_pos == (3, 2, 0) + + +if __name__ == "__main__": + pytest.main([__file__, "-q"]) diff --git a/tests/test_tropical_counting.py b/tests/test_tropical_counting.py new file mode 100644 index 00000000..84bbce19 --- /dev/null +++ b/tests/test_tropical_counting.py @@ -0,0 +1,364 @@ +import numpy as np +import tensorcircuit as tc +from applications.tropical_algebra import ( + _counting_tensordot, + split_energy_count, +) + + +def _be(): + return tc.backend + + +def _ref_counting_tensordot(anp, bnp, axes, eps=1e-9): + # brute-force (max, degeneracy) tensordot on energy-only inputs (count starts at 1) + if isinstance(axes, int): + a_axes = list(range(anp.ndim - axes, anp.ndim)) + b_axes = list(range(0, axes)) + else: + a_axes, b_axes = list(axes[0]), list(axes[1]) + a_free = [i for i in range(anp.ndim) if i not in a_axes] + b_free = [i for i in range(bnp.ndim) if i not in b_axes] + at = np.transpose(anp, a_free + a_axes) + bt = np.transpose(bnp, b_axes + b_free) + a_fs = [anp.shape[i] for i in a_free] + b_fs = [bnp.shape[i] for i in b_free] + k = int(np.prod([anp.shape[i] for i in a_axes]) or 1) + A = int(np.prod(a_fs) or 1) + b_n = int(np.prod(b_fs) or 1) + a2 = at.reshape(A, k) + b2 = bt.reshape(k, b_n) + out_e = np.full((A, b_n), -np.inf) + out_n = np.zeros((A, b_n)) + for i in range(A): + for j in range(b_n): + s = a2[i, :] + b2[:, j] + mx = np.max(s) + out_e[i, j] = mx + out_n[i, j] = int(np.sum(np.abs(s - mx) < eps)) # degeneracy = #max-tied + return out_e.reshape(a_fs + b_fs), out_n.reshape(a_fs + b_fs) + + +def test_counting_tensordot_matrix(): + be = _be() + rng = np.random.default_rng(10) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + # stacked inputs: [x, n], counts init to 1 + a = be.cast( + be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], axis=-1)), "float64" + ) + b = be.cast( + be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], axis=-1)), "float64" + ) + got = np.array(_counting_tensordot(be, a, b, axes=1)) + ref_e, ref_n = _ref_counting_tensordot(anp, bnp, 1) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_split_energy_count(): + be = _be() + arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + stacked = be.cast( + be.convert_to_tensor(np.stack([arr, arr * 2], axis=-1)), "float64" + ) + e, n = split_energy_count(stacked) + np.testing.assert_allclose(np.array(e), arr) + np.testing.assert_allclose(np.array(n), arr * 2) + + +# --- Task B2: counting einsum (hyperedge / copy-node aware) --- + +from applications.tropical_algebra import ( # noqa: E402 + _counting_einsum, +) + + +def _ref_counting_einsum(eq, a, b, eps=1e-9): + """Brute-force (energy, degeneracy) 2-operand einsum on energy-only inputs.""" + import itertools + + lhs, rhs = eq.split("->") + ia, ib = lhs.split(",") + sizes = {} + for s, t in zip([ia, ib], [a, b]): + for c, dim in zip(s, t.shape): + sizes[c] = dim + out_e = np.full([sizes[c] for c in rhs], -np.inf) + out_n = np.zeros([sizes[c] for c in rhs]) + allc = list(dict.fromkeys(list(ia) + list(ib))) + for combo in itertools.product(*[range(sizes[c]) for c in allc]): + env = dict(zip(allc, combo)) + e = a[tuple(env[c] for c in ia)] + b[tuple(env[c] for c in ib)] + oidx = tuple(env[c] for c in rhs) + if e > out_e[oidx] + eps: + out_e[oidx] = e + out_n[oidx] = 1 + elif abs(e - out_e[oidx]) < eps: + out_n[oidx] += 1 + return out_e, out_n + + +def test_counting_einsum_pair(): + be = _be() + rng = np.random.default_rng(11) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,bc->ac", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,bc->ac", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_counting_einsum_hyperedge(): + # shared index 'a' (batch) — copy-node pair shape + be = _be() + rng = np.random.default_rng(12) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,ac->abc", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,ac->abc", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_counting_einsum_tie_degeneracy(): + # Symmetric inputs over the contracted axis -> a clean tie (degeneracy > 1). + be = _be() + # a: shape (1, 2) with equal values on the contracted axis; b: shape (2, 1). + anp = np.array([[1.0, 1.0]]) # both k=0 and k=1 contribute the same energy + bnp = np.array([[2.0], [2.0]]) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,bc->ac", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,bc->ac", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + # explicit assertion: energy 3.0 achieved by both k values -> degeneracy 2 + assert ref_n[0, 0] == 2 + np.testing.assert_allclose(got[0, 0, 0], 3.0) + np.testing.assert_allclose(got[0, 0, 1], 2.0) + + +def test_counting_einsum_single_operand_diagonal(): + # eq "aa->a": diagonal, then no contraction -> energy=diag(e), count=diag(n) + import tensorcircuit.backends.numpy_backend as nb + + be = nb.NumpyBackend() + e = np.array([[1.0, 5.0], [3.0, 2.0]]) + n = np.array([[1.0, 2.0], [1.0, 1.0]]) + a = np.stack([e, n], axis=-1) + out = _counting_einsum(be, "aa->a", a) + np.testing.assert_allclose(out[..., 0], np.diag(e)) + np.testing.assert_allclose(out[..., 1], np.diag(n)) + + +def test_counting_einsum_single_operand_reduce(): + # eq "ab->a": max over b of energy, tie-sum count + import tensorcircuit.backends.numpy_backend as nb + + be = nb.NumpyBackend() + e = np.array([[1.0, 5.0], [3.0, 3.0]]) + n = np.array([[1.0, 2.0], [1.0, 4.0]]) + a = np.stack([e, n], axis=-1) + out = _counting_einsum(be, "ab->a", a) + np.testing.assert_allclose(out[..., 0], [5.0, 3.0]) + np.testing.assert_allclose(out[..., 1], [2.0, 5.0]) # ties: b=0,1 both 3 -> 1+4 + + +def test_counting_einsum_two_operand_intra_repeat(): + # eq "aa,b->ab": operand a has repeated index + import tensorcircuit.backends.numpy_backend as nb + + be = nb.NumpyBackend() + e_a = np.array([[2.0, 1.0], [1.0, 2.0]]) + n_a = np.ones((2, 2)) + a = np.stack([e_a, n_a], axis=-1) + e_b = np.array([0.0, 0.0]) + n_b = np.array([1.0, 1.0]) + b = np.stack([e_b, n_b], axis=-1) + out = _counting_einsum(be, "aa,b->ab", a, b) + # diagonal of a is [2,2]; outer with b=[0,0] -> [[2,2],[2,2]], counts follow + np.testing.assert_allclose(out[..., 0], [[2, 2], [2, 2]]) + # count stream: diag(n_a)=[1,1] outer n_b=[1,1] -> [[1,1],[1,1]] + np.testing.assert_allclose(out[..., 1], [[1, 1], [1, 1]]) + + +def test_counting_einsum_multi_axis(): + # Multi-axis contraction "ab,ac->c": shared index a (contracted) plus a free + # index b of operand a (also summed out) -> two axes reduced. Compared to a + # brute-force (energy, degeneracy) reference. + be = _be() + rng = np.random.default_rng(14) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,ac->c", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,ac->c", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +# --- Task 8: degeneracy side-channel + standard-contraction clear discipline --- + + +def test_degeneracy_none_after_standard_contraction(): + """A counting contraction stashes count=deg into the aux side-channel; a + subsequent STANDARD contraction must clear it (unconditional clear in + ``cons._algebraic_base_contraction``), so ``degeneracy()`` returns None + rather than a stale count from the earlier counting contraction. + """ + import opt_einsum + import tensornetwork as tn + import tensorcircuit.cons as cons + import applications.tropical_algebra as tr + from tests._tropical_test_utils import ( + build_ising_tn, + brute_force_energy_and_degeneracy, + ) + + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] # 5-ring + J = [1.0, 1.0, 1.0, 1.0, 1.0] # ferromagnetic + h = [0.0, 0.0, 0.0, 0.0, 0.0] # zero field -> Z2 spin-flip symmetry, g=2 + expected_e, expected_n = brute_force_energy_and_degeneracy(len(spins), edges, J, h) + assert ( + expected_n == 2 + ) # guard: this instance MUST be degenerate or the test is moot + nodes = build_ising_tn(spins, edges, J, h) + + with tr.counting_tropical(): + node = cons._algebraic_base_contraction( + nodes, + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[], + ignore_edge_order=True, + ) + # counting decode stashed the count; degeneracy() reads it (valid here) + np.testing.assert_allclose(np.array(tr.degeneracy()), expected_n) + np.testing.assert_allclose(np.array(node.tensor), -expected_e, atol=1e-6) + + # exit -> StandardAlgebra restored. A standard contraction must clear aux + # (the unconditional _stash_aux_outputs({}) at the top of every contraction), + # so a stale count cannot leak out of the counting block. + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal((2, 3)).astype(np.complex64)) + b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) + tn.connect(a[1], b[0]) + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], b[1]], + ) + assert tr.degeneracy() is None # standard contraction cleared aux + + +# --- Task 10: non-scalar counting canary (aux co-indexed with output_edge_order) --- + + +def test_nonscalar_counting_energy_and_degeneracy_per_output(): + """A counting contraction with ONE free (dangling) spin returns energy AND + degeneracy PER output configuration, with the count aux co-indexed with the + energy via ``output_edge_order``. + + Validates two things at once: + (1) non-scalar counting works at all -- ``decode`` strips the trailing count + axis so ``tn.Node`` wraps a rank-correct energy tensor (rank == #free); + (2) the aux-reorder logic (Task 5's ``output_edge_order`` permutation applied + to aux) aligns count with energy. If the perm were wrong, count would be + permuted differently from energy and the per-output assertion would fail. + + Brute force: for each value v in {0,1} of the free spin (cfg=0 -> s=+1, + cfg=1 -> s=-1), enumerate the other spins, find min E and its degeneracy. + """ + import opt_einsum + import applications.tropical_algebra as tr + import tensorcircuit.cons as cons + from tests._tropical_test_utils import ( + build_ring_with_free_spin, + brute_nonscalar_counting, + ) + + nodes, free_edge, J, h = build_ring_with_free_spin() + expected_e, expected_n = brute_nonscalar_counting(nodes, free_edge, J, h) + + with tr.counting_tropical(): + node = cons._algebraic_base_contraction( + nodes, + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[free_edge], + ) + E = np.asarray(node.tensor) + N = np.asarray(tr.degeneracy()) + + np.testing.assert_allclose(E, expected_e) + np.testing.assert_allclose(N, expected_n) + # guard: the two output configs must differ in BOTH energy and count, + # otherwise a permuted aux could still pass the assertion by accident. + assert E[0] != E[1] or N[0] != N[1] + + +# --- Task 10 follow-up: 2-free-spin perm-direction canary --- + + +def test_nonscalar_counting_two_free_spins_reversed_order(): + """A counting contraction with TWO free (dangling) spins where + ``output_edge_order`` is INTENTIONALLY REVERSED relative to the dangling-edge + order, so the aux-reorder permutation in ``cons._algebraic_base_contraction`` + (``perm = [dangling_edges.index(e) for e in order]``) is non-trivial + (``(1, 0)`` swap, not the identity ``(0,)`` exercised by the one-free-spin + canary). + + Validates that the SAME permutation applied to the energy tensor (via + ``final_node.reorder_edges``) is also applied to the aux degeneracy tensor + (via ``kbe.transpose(v, tuple(perm))``). If the aux perm were ELIDED, the + energy ``E`` would come out in the requested ``[b, a]`` order but the count + ``N`` would be left in the dangling ``[a, b]`` order, so + ``assert_allclose(N, expected_n_ab.T)`` would FAIL (the brute-force result is + asymmetric under transpose by parameter choice -- see the guard below). + + The energy ``E`` itself is also checked against the transposed brute force, + pinning down which order the contraction actually produced. + """ + import opt_einsum + import applications.tropical_algebra as tr + import tensorcircuit.cons as cons + from tests._tropical_test_utils import ( + build_ring_with_two_free_spins, + brute_nonscalar_counting_2d, + ) + + nodes, edge_a, edge_b, J, h = build_ring_with_two_free_spins() + # brute force in the natural (a, b) order: + expected_e_ab, expected_n_ab = brute_nonscalar_counting_2d( + nodes, edge_a, edge_b, J, h + ) + # guard: the brute-force result MUST be asymmetric under transpose, otherwise + # a wrong-order aux could still pass ``assert_allclose(N, expected_n_ab.T)`` + # by accident. The default parameters give E=[[5,7],[5,-1]], N=[[1,2],[1,1]], + # both asymmetric. + assert not np.allclose(expected_e_ab, expected_e_ab.T), ( + "weak canary: brute-force energy is symmetric under transpose, so the " + "perm direction is not load-bearing for E" + ) + assert not np.allclose(expected_n_ab, expected_n_ab.T), ( + "weak canary: brute-force degeneracy is symmetric under transpose, so " + "the perm direction is not load-bearing for N" + ) + + # contract with output_edge_order REVERSED -> [edge_b, edge_a]; under + # sorted_edges the dangling order is [edge_a, edge_b], so the perm is (1, 0). + with tr.counting_tropical(): + node = cons._algebraic_base_contraction( + nodes, + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[edge_b, edge_a], + ) + E = np.asarray(node.tensor) # in [b, a] order (reorder_edges applied) + N = np.asarray(tr.degeneracy()) # MUST also be in [b, a] order if the + # # aux perm is applied correctly. + # expected in [b, a] order = transpose of the (a, b) brute force: + np.testing.assert_allclose(E, expected_e_ab.T) + np.testing.assert_allclose(N, expected_n_ab.T) diff --git a/tests/test_tropical_example.py b/tests/test_tropical_example.py new file mode 100644 index 00000000..40507345 --- /dev/null +++ b/tests/test_tropical_example.py @@ -0,0 +1,39 @@ +import numpy as np +import tensornetwork as tn +import tensorcircuit as tc +import tensorcircuit.cons as cons +from applications.tropical_algebra import ( + MaxPlusAlgebra, + tropical, + recover_configuration, +) + + +def test_example_maxplus_tensordot_matches_brute(): + be = tc.backend + rng = np.random.default_rng(0) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(MaxPlusAlgebra().tensordot(be, a, b, 1)) + # brute max-plus: Y[i,j] = max_k a[i,k] + b[k,j] + ref = np.max(anp[:, :, None] + bnp[None, :, :], axis=1) + np.testing.assert_allclose(got, ref) + + +def test_example_recover_configuration_on_tiny_ising(): + # 3-node chain (vector-matrix-vector) -> recover_configuration round-trips + be = tc.backend + n = tn.Node(be.cast(be.convert_to_tensor(np.array([0.0, 0.0])), "float64")) + e = tn.Node( + be.cast(be.convert_to_tensor(np.array([[1.0, -1.0], [-1.0, 1.0]])), "float64") + ) + m = tn.Node(be.cast(be.convert_to_tensor(np.array([0.5, -0.5])), "float64")) + tn.connect(n[0], e[0]) + tn.connect(e[1], m[0]) + nodes = [n, e, m] + with tropical(track=True): + float(cons.contractor(nodes, output_edge_order=[]).tensor) + cfg = recover_configuration() + assert isinstance(cfg, dict) + assert len(cfg) >= 1 diff --git a/tests/test_tropical_ising.py b/tests/test_tropical_ising.py new file mode 100644 index 00000000..14ced604 --- /dev/null +++ b/tests/test_tropical_ising.py @@ -0,0 +1,223 @@ +import itertools +import numpy as np +import tensorcircuit.cons as cons +from applications.tropical_algebra import tropical, counting_tropical +from tests._tropical_test_utils import ( + build_ising_tn, + brute_force_energy, + brute_force_energy_and_degeneracy, +) + + +def test_ising_ring_ground_state_matches_bruteforce(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] # 5-ring + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + nodes = build_ising_tn(spins, edges, J, h) + + e_ground = brute_force_energy(spins, edges, J, h) + + with tropical(): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + + # contraction returns max_cfg(-E) = -E_ground + np.testing.assert_allclose(val, -e_ground, atol=1e-6) + + +# --- Task A2: configuration recovery (argmax backtracking) canary --- + + +def _ising_energy_of_cfg(cfg, edges, j_vals, h): + """Ising energy of a {0,1}-config: E = -sum J s_i s_j - sum h s_i (s=1-2cfg).""" + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + si, sj = 1 - 2 * cfg[i], 1 - 2 * cfg[j] + e -= jij * si * sj + for i, hi in zip(range(len(h)), h): + e -= hi * (1 - 2 * cfg[i]) + return e + + +def _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors): + """Max-plus total of a symbol->value assignment = sum of tensor entries = + -E for that configuration. Mapping-independent optimality proof.""" + total = 0.0 + for term, t in zip(input_sets, raw_tensors): + idx = tuple(int(cfg_sym[c]) for c in term) + total += float(np.asarray(t)[idx]) + return total + + +def test_ising_config_recovery_five_ring(): + """Task A2 canary: recover an optimal spin configuration of the 5-spin Ising + ring under ``tropical(track=True)`` and verify its energy equals the + brute-force ground energy. + + The 5-ring mixes ``tensordot`` (ordinary pair) and ``einsum`` (hyperedge -- + a CopyNode hub of degree >= 2 puts its symbol on >= 3 regular nodes) + contraction steps, so this exercises both backtracking branches. Optimality + is asserted two ways: (1) the max-plus total of the recovered symbol + assignment equals the contraction value (== -E_ground) -- mapping- + independent; (2) translating the assignment to per-spin values via the + rank-1 field tensors and computing the Ising energy gives ``best_e``. + """ + from applications.tropical_algebra import ( + get_recorded_topology, + recover_configuration, + ) + + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] # 5-ring + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + + # brute-force ground energy over the 2^5 configs + best_e = None + for cfg in itertools.product([0, 1], repeat=len(spins)): + e = _ising_energy_of_cfg(cfg, edges, J, h) + if best_e is None or e < best_e - 1e-9: + best_e = e + + nodes = build_ising_tn(spins, edges, J, h) + with tropical(track=True): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + cfg_sym = recover_configuration() + + # contraction value == -E_ground + np.testing.assert_allclose(val, -best_e, atol=1e-6) + + _tree, input_sets, raw_tensors = get_recorded_topology() + assert input_sets is not None and raw_tensors is not None + # every index label received a value + all_labels = set().union(*[set(t) for t in input_sets]) + assert set(cfg_sym.keys()) == all_labels + + # (1) mapping-independent optimality: max-plus total == contraction value + recovered_total = _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, val, atol=1e-6) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + # (2) translate symbol config -> per-spin cfg via rank-1 field tensors + # (Tv_i has a single index label; its value is cfg[i]). The field tensors + # appear in input_sets in spin order (regular_nodes sorted by _stable_id_, + # which follows creation order in build_ising_tn). + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg_sym[term[0]]) + assert len(spin_cfg) == len(spins) + recovered_e = _ising_energy_of_cfg( + [spin_cfg[i] for i in range(len(spins))], edges, J, h + ) + np.testing.assert_allclose(recovered_e, best_e, atol=1e-6) + + +def test_ising_config_recovery_degenerate_ring(): + """Config recovery under degeneracy (Z2-symmetric ferromagnetic ring, zero + field): the recovered config must still be energy-optimal (it is one of the + degenerate ground states, selected by first-argument-wins tie-breaking).""" + from applications.tropical_algebra import ( + get_recorded_topology, + recover_configuration, + ) + + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, 1.0, 1.0, 1.0, 1.0] # ferromagnetic + h = [0.0, 0.0, 0.0, 0.0, 0.0] # zero field -> Z2 spin-flip symmetry, g=2 + + best_e = None + for cfg in itertools.product([0, 1], repeat=len(spins)): + e = _ising_energy_of_cfg(cfg, edges, J, h) + if best_e is None or e < best_e - 1e-9: + best_e = e + assert best_e == -5.0 # ferromagnetic ground state + + nodes = build_ising_tn(spins, edges, J, h) + with tropical(track=True): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + cfg_sym = recover_configuration() + + np.testing.assert_allclose(val, -best_e, atol=1e-6) + _tree, input_sets, raw_tensors = get_recorded_topology() + recovered_total = _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + # the recovered spin config must be one of the two aligned ground states + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg_sym[term[0]]) + recovered_e = _ising_energy_of_cfg( + [spin_cfg[i] for i in range(len(spins))], edges, J, h + ) + np.testing.assert_allclose(recovered_e, best_e, atol=1e-6) + + +# --- Task B3: counting tropical (energy*, degeneracy) canary --- + + +def test_ising_counting_ground_state_and_degeneracy(): + """Energy via ``node.tensor``; degeneracy via ``tr.degeneracy()`` side channel. + + Under the encode/decode design, ``CountingRepresentation.encode`` attaches + count=1 to each leaf, so the test builder emits PLAIN energy tensors (the + same ``build_ising_tn`` used by the max-plus tests). After contraction under + ``counting_tropical()``, the primary tensor is the energy (-E_ground) and + the count is stashed in the aux side-channel for ``degeneracy()`` to read. + """ + import applications.tropical_algebra as tr + + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] # 5-ring + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + # brute-force ground energy + degeneracy (s = 1 - 2*cfg; cfg=0 -> s=+1, cfg=1 -> s=-1) + best_e, deg = brute_force_energy_and_degeneracy(len(spins), edges, J, h) + + nodes = build_ising_tn(spins, edges, J, h) + + with counting_tropical(): + res = cons.contractor(nodes, output_edge_order=[], ignore_edge_order=True) + count = tr.degeneracy() + + energy = np.array(res.tensor) + # contraction returns max_cfg(-E) = -E_ground and the #cfgs achieving it. + np.testing.assert_allclose(energy, -best_e, atol=1e-6) + np.testing.assert_allclose(np.array(count), deg, atol=1e-6) + + +def test_ising_counting_degenerate_ring(): + """Degenerate ground state (g=2) canary for the counting tropical stream. + + The original ``test_ising_counting_ground_state_and_degeneracy`` instance is + tuned to a *unique* ground state (g=1), so a regression that always returns + ``count=1`` would pass undetected. A ferromagnetic ring with zero field has a + Z2 (global spin-flip) symmetry -> exactly 2 ground states (all-up and + all-down), so this case guards the degeneracy stream against such a regression. + """ + import applications.tropical_algebra as tr + + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] # 5-ring + J = [1.0, 1.0, 1.0, 1.0, 1.0] # ferromagnetic + h = [0.0, 0.0, 0.0, 0.0, 0.0] # zero field -> Z2 spin-flip symmetry + # brute-force ground energy + degeneracy (s = 1 - 2*cfg; cfg=0 -> s=+1, cfg=1 -> s=-1) + best_e, deg = brute_force_energy_and_degeneracy(len(spins), edges, J, h) + + # sanity-check the derivation: 5 ferromagnetic bonds, all J=+1, h=0 -> + # ground state all-spins-aligned, E_ground = -sum(J) = -5 over exactly 2 + # configs (Z2-related). Contraction returns max_cfg(-E) = -E_ground = 5. + assert best_e == -5.0 + assert deg == 2 + + nodes = build_ising_tn(spins, edges, J, h) + + with counting_tropical(): + res = cons.contractor(nodes, output_edge_order=[], ignore_edge_order=True) + count = tr.degeneracy() + + energy = np.array(res.tensor) + np.testing.assert_allclose(energy, -best_e, atol=1e-6) # == 5 + np.testing.assert_allclose(np.array(count), deg, atol=1e-6) # == 2 From a1f8c1b4b338cf026c35daaf7806311c52e8a448 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sat, 18 Jul 2026 20:42:30 +0800 Subject: [PATCH 02/16] bf16: genuine-bf16 einsum kernel --- applications/bcomplex32_algebra.py | 118 ++++++++++++++++++++++++----- tests/test_bcomplex32_algebra.py | 39 ++++++++++ 2 files changed, 139 insertions(+), 18 deletions(-) diff --git a/applications/bcomplex32_algebra.py b/applications/bcomplex32_algebra.py index 5d410124..43d1b014 100644 --- a/applications/bcomplex32_algebra.py +++ b/applications/bcomplex32_algebra.py @@ -47,28 +47,110 @@ def _pair_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: def _pair_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: - """Complex einsum = real bf16 einsums (4M for 2 operands, 2 for 1). - - Each pair component is upcast bf16 -> float32 before ``be.einsum``: numpy's - einsum does not support ``ml_dtypes.bfloat16`` (``TypeError: invalid data - type for einsum``), and float32 is the universal portable compute dtype for - the 4M decomposition. The bf16 quantization happened at encode time - (``_complex_to_pair``), so the values flowing in are already bf16-quantized; - this cast only widens the compute dtype, it does not undo quantization. - ``_pair_tensordot`` needs no such cast because ``np.tensordot`` accepts bf16. + """Complex einsum = 4 real bf16 einsums (4M for 2 operands, 2 for 1). + + **Two-operand (genuine bf16):** manually decomposed into ``be.tensordot`` + + ``be.transpose``, because numpy's C ``einsum`` rejects ``ml_dtypes.bfloat16`` + (it is a standalone C routine with a hardcoded dtype allowlist, not a + ufunc). ``np.tensordot`` accepts bf16 because it dispatches through ufunc + loops, so the compute is genuine bf16 end-to-end — no float32 upcast. + The decomposition parses the einsum subscript equation to find contracted + axes, then uses ``tensordot`` for the contraction and ``transpose`` to + match the output subscript order. + + **Single-operand (genuine bf16):** decomposed into ``np.diagonal`` + + ``be.sum`` + ``be.transpose``, all bf16-safe. Handles reductions, + transposes, diagonals, and traces — the full einsum single-operand + semantics without any float32 upcast. + + ``_pair_tensordot`` needs no such routing because ``np.tensordot`` already + preserves bf16 (verified: it accumulates in bf16, not float32). + cotengra feeds only 1-2-operand equations here (it decomposes hyperedges + itself), all of which are pairwise contractions that map cleanly to + tensordot. """ if len(operands) == 1: + # ── Single-operand: pure bf16 (decompose into sum + transpose + diagonal) ── a = operands[0] - ar = be.cast(a[..., 0], "float32") - ai = be.cast(a[..., 1], "float32") - return be.stack([be.einsum(eq, ar), be.einsum(eq, ai)], axis=-1) + + # Implicit mode (no ``->``) is an identity / no-op at the einsum level. + if "->" not in eq: + return a + + lhs, out_subs = eq.split("->") + + def _single_op(x: Tensor) -> Tensor: + """Apply a 1-operand einsum to one bf16 half, staying in bf16.""" + x_subs = list(lhs) # mutable subscript list we update in place + + # Step 1 — diagonalise every repeated index. + # ``np.diagonal(x, axis1=p0, axis2=p1)`` removes axis-p1, then + # appends the diagonal (size = common dim) at the end. + while True: + dup = next((c for c in set(x_subs) if x_subs.count(c) > 1), None) + if dup is None: + break + pos = [i for i, c in enumerate(x_subs) if c == dup] + x = np.diagonal(x, axis1=pos[0], axis2=pos[-1]) + x_subs = [c for i, c in enumerate(x_subs) if i != pos[-1]] + [dup] + + # Step 2 — sum over indices NOT wanted in the output. + out_set = set(out_subs) + sum_indices = [c for c in x_subs if c not in out_set] + if sum_indices: + x = be.sum(x, axis=tuple(x_subs.index(c) for c in sum_indices)) + x_subs = [c for c in x_subs if c in out_set] + + # Step 3 — transpose remaining indices into the requested output order. + if x_subs != list(out_subs): + perm = tuple(x_subs.index(c) for c in out_subs) + x = be.transpose(x, perm) + + return x + + return be.stack( + [_single_op(a[..., 0]), _single_op(a[..., 1])], + axis=-1, + ) + + # ── 2-operand: bf16-safe tensordot decomposition ────────────────── a, b = operands - ar = be.cast(a[..., 0], "float32") - ai = be.cast(a[..., 1], "float32") - br = be.cast(b[..., 0], "float32") - bi = be.cast(b[..., 1], "float32") - cr = be.einsum(eq, ar, br) - be.einsum(eq, ai, bi) - ci = be.einsum(eq, ar, bi) + be.einsum(eq, ai, br) + ar, ai = a[..., 0], a[..., 1] + br, bi = b[..., 0], b[..., 1] + + # Parse the einsum equation once + if "->" not in eq: + raise ValueError( + f"implicit-mode einsum {eq!r} not supported for bf16; " f"use explicit '->'" + ) + lhs, out_subs = eq.split("->") + a_subs, b_subs = lhs.split(",") + + a_set: set[str] = set(a_subs) + b_set: set[str] = set(b_subs) + contracted = [c for c in a_subs if c in b_set] + + a_free = [c for c in a_subs if c not in b_set] + b_free = [c for c in b_subs if c not in a_set] + out_order = list(out_subs) + + def _contract(x: Tensor, y: Tensor) -> Tensor: + """Pairwise bf16-safe einsum → tensordot + optional transpose.""" + if contracted: + a_axes = [a_subs.index(c) for c in contracted] + b_axes = [b_subs.index(c) for c in contracted] + result = be.tensordot(x, y, axes=(a_axes, b_axes)) + else: + result = be.tensordot(x, y, axes=0) + + free_order = a_free + b_free + if free_order != out_order: + perm = [free_order.index(c) for c in out_order] + result = be.transpose(result, perm) + return result + + cr = _contract(ar, br) - _contract(ai, bi) + ci = _contract(ar, bi) + _contract(ai, br) return be.stack([cr, ci], axis=-1) diff --git a/tests/test_bcomplex32_algebra.py b/tests/test_bcomplex32_algebra.py index 3c584712..d182d9de 100644 --- a/tests/test_bcomplex32_algebra.py +++ b/tests/test_bcomplex32_algebra.py @@ -88,3 +88,42 @@ def test_bf16_wall_avoidance_canary(): c2.H(0) c2.cnot(0, 1) # subsequent native contraction assert np.asarray(c2.state()).shape == (4,) # algebra restored, no leak + + +def test_pair_einsum_keeps_bfloat16_dtype(): + """T1: _pair_einsum must compute in bf16, not upcast to float32. + + numpy's np.einsum rejects bf16; the old _pair_einsum worked around it by + upcasting to float32 (so its output pair was float32, not bf16). The rewrite + uses manual tensordot decomposition, which stays bf16. This test locks that. + """ + import ml_dtypes + + bf = ml_dtypes.bfloat16 + a = np.array([[1.0 + 2.0j, 3.0j], [-1.0j, 2.0 - 1.0j]], dtype=np.complex64) + b = np.array([[0.5 + 0.5j, 1.0j], [2.0j, -1.0 + 1.0j]], dtype=np.complex64) + pa = _complex_to_pair(be, a) # bf16 pair, shape (2, 2, 2) + pb = _complex_to_pair(be, b) + out = np.asarray(_pair_einsum(be, "ij,jk->ik", pa, pb)) + assert out.dtype == bf, f"_pair_einsum upcast to {out.dtype}; expected bfloat16" + + +def test_bf16_ghz8_runs_and_matches_native(): + """T3: the 8-qubit GHZ that previously crashed (cotengra autoray transpose on + a pair result) now runs under bcomplex32 and matches native within bf16 + tolerance. prefer_einsum=True avoids the transpose path; the genuine-bf16 + kernel keeps intermediates bf16 end-to-end. + """ + + def ghz(n): + c = tc.Circuit(n) + c.H(0) + for i in range(n - 1): + c.cnot(i, i + 1) + return np.asarray(c.state()) + + ref = ghz(8) + with bcomplex32(): + got = ghz(8) + assert got.shape == ref.shape + assert np.allclose(got, ref, rtol=5e-2), f"max abs diff = {np.abs(got - ref).max()}" From b3d3d345a207b62a0ccffb65055c0b6b95d3ad61 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sat, 18 Jul 2026 22:20:24 +0800 Subject: [PATCH 03/16] fix: reduce cognitive complexity in _pair_einsum and fix implicit string concatenation --- applications/bcomplex32_algebra.py | 72 +++++++++++++++++------------- 1 file changed, 40 insertions(+), 32 deletions(-) diff --git a/applications/bcomplex32_algebra.py b/applications/bcomplex32_algebra.py index 43d1b014..82798515 100644 --- a/applications/bcomplex32_algebra.py +++ b/applications/bcomplex32_algebra.py @@ -9,6 +9,8 @@ from typing import Any, Dict, Iterator, List, Tuple import contextlib +import numpy as np + import tensorcircuit.cons as cons from tensorcircuit.contraction_algebra import ContractionAlgebra, Representation @@ -46,6 +48,39 @@ def _pair_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: return be.stack([cr, ci], axis=-1) +def _einsum_single_operand_half( + be: Backend, x: Tensor, lhs: str, out_subs: str +) -> Tensor: + """Apply a 1-operand einsum to one bf16 half (decomposed into diagonal + sum + + transpose), staying in bf16 end-to-end. Handles reductions, transposes, + diagonals, and traces without float32 upcast. + """ + x_subs = list(lhs) # mutable subscript list we update in place + + # Step 1 — diagonalise every repeated index. + while True: + dup = next((c for c in set(x_subs) if x_subs.count(c) > 1), None) + if dup is None: + break + pos = [i for i, c in enumerate(x_subs) if c == dup] + x = np.diagonal(x, axis1=pos[0], axis2=pos[-1]) + x_subs = [c for i, c in enumerate(x_subs) if i != pos[-1]] + [dup] + + # Step 2 — sum over indices NOT wanted in the output. + out_set = set(out_subs) + sum_indices = [c for c in x_subs if c not in out_set] + if sum_indices: + x = be.sum(x, axis=tuple(x_subs.index(c) for c in sum_indices)) + x_subs = [c for c in x_subs if c in out_set] + + # Step 3 — transpose remaining indices into the requested output order. + if x_subs != list(out_subs): + perm = tuple(x_subs.index(c) for c in out_subs) + x = be.transpose(x, perm) + + return x + + def _pair_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: """Complex einsum = 4 real bf16 einsums (4M for 2 operands, 2 for 1). @@ -78,38 +113,11 @@ def _pair_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: return a lhs, out_subs = eq.split("->") - - def _single_op(x: Tensor) -> Tensor: - """Apply a 1-operand einsum to one bf16 half, staying in bf16.""" - x_subs = list(lhs) # mutable subscript list we update in place - - # Step 1 — diagonalise every repeated index. - # ``np.diagonal(x, axis1=p0, axis2=p1)`` removes axis-p1, then - # appends the diagonal (size = common dim) at the end. - while True: - dup = next((c for c in set(x_subs) if x_subs.count(c) > 1), None) - if dup is None: - break - pos = [i for i, c in enumerate(x_subs) if c == dup] - x = np.diagonal(x, axis1=pos[0], axis2=pos[-1]) - x_subs = [c for i, c in enumerate(x_subs) if i != pos[-1]] + [dup] - - # Step 2 — sum over indices NOT wanted in the output. - out_set = set(out_subs) - sum_indices = [c for c in x_subs if c not in out_set] - if sum_indices: - x = be.sum(x, axis=tuple(x_subs.index(c) for c in sum_indices)) - x_subs = [c for c in x_subs if c in out_set] - - # Step 3 — transpose remaining indices into the requested output order. - if x_subs != list(out_subs): - perm = tuple(x_subs.index(c) for c in out_subs) - x = be.transpose(x, perm) - - return x - return be.stack( - [_single_op(a[..., 0]), _single_op(a[..., 1])], + [ + _einsum_single_operand_half(be, a[..., 0], lhs, out_subs), + _einsum_single_operand_half(be, a[..., 1], lhs, out_subs), + ], axis=-1, ) @@ -121,7 +129,7 @@ def _single_op(x: Tensor) -> Tensor: # Parse the einsum equation once if "->" not in eq: raise ValueError( - f"implicit-mode einsum {eq!r} not supported for bf16; " f"use explicit '->'" + f"implicit-mode einsum {eq!r} not supported for bf16; use explicit '->'" ) lhs, out_subs = eq.split("->") a_subs, b_subs = lhs.split(",") From 5de41f8af9a6e501480b54a8cb3a4253963adf95 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sat, 18 Jul 2026 23:46:11 +0800 Subject: [PATCH 04/16] refactor: merge 6 tropical test files into single test_tropical.py --- tests/_tropical_test_utils.py | 379 ---------- tests/test_tropical.py | 1139 +++++++++++++++++++++++++++++++ tests/test_tropical_algebra.py | 272 -------- tests/test_tropical_config.py | 246 ------- tests/test_tropical_counting.py | 364 ---------- tests/test_tropical_example.py | 39 -- tests/test_tropical_ising.py | 223 ------ 7 files changed, 1139 insertions(+), 1523 deletions(-) delete mode 100644 tests/_tropical_test_utils.py create mode 100644 tests/test_tropical.py delete mode 100644 tests/test_tropical_algebra.py delete mode 100644 tests/test_tropical_config.py delete mode 100644 tests/test_tropical_counting.py delete mode 100644 tests/test_tropical_example.py delete mode 100644 tests/test_tropical_ising.py diff --git a/tests/_tropical_test_utils.py b/tests/_tropical_test_utils.py deleted file mode 100644 index 9620e73f..00000000 --- a/tests/_tropical_test_utils.py +++ /dev/null @@ -1,379 +0,0 @@ -"""Shared helpers for the tropical Ising tensor-network tests. - -Extracted from ``tests/test_tropical_ising.py`` and -``tests/test_tropical_config.py`` to remove duplication between their -near-identical CopyNode Ising TN builders and brute-force energy -enumeration loops. -""" - -import itertools - -import numpy as np -import tensornetwork as tn - -import tensorcircuit as tc - - -def build_ising_tn(spins, edges, j_vals, h): - """Build the tropical (max-plus) Ising tensor network as tensornetwork nodes. - - Per spin: a ``tn.CopyNode`` of rank (degree+1), dimension 2. A CopyNode is a true - hyperedge hub -- tensornetwork/cotengra identify all its legs as one shared index - (the delta/copy constraint is structural, not stored as data). This is what makes - the contraction route hyperedge-containing pairs through the tropical **einsum** - branch and ordinary pairs through the tropical **tensordot** branch, covering both. - Per edge (i,j): Te[si,sj] = J*si*sj -> [[J,-J],[-J,J]]. - Per spin: field Tv[s] = h*s -> [h,-h]. - Contracting over max-plus yields max_cfg(-E) = -E_ground. - - Note: an earlier draft built the delta as a dense ``tn.Node`` (0 on diagonal, a NEG - proxy off-diagonal). That is a *regular* node, not a hyperedge, so it does not - trigger the einsum branch and -- worse -- ``cons.contractor``'s ``_merge_single_gates`` - preprocessing collapses the whole network to one node before the algebraic path - runs, silently bypassing the tropical primitives (result 0.0). Using ``tn.CopyNode`` - is both the faithful "copy node -> hyperedge" construction and the one that works - end-to-end through the public ``cons.contractor`` API. - """ - be = tc.backend - degree = dict.fromkeys(spins, 0) - for i, j in edges: - degree[i] += 1 - degree[j] += 1 - - nodes = [] - copy_nodes = {} - # one field leg + one leg per incident edge - for i in spins: - cn = tn.CopyNode(degree[i] + 1, 2) # rank (degree+1) delta, dim 2 -> hyperedge - copy_nodes[i] = cn - nodes.append(cn) - # field tensor on the first leg - tv = np.array([h[i], -h[i]], dtype=np.float64) - tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) - tn.connect(cn[0], tvn[0]) - nodes.append(tvn) - # bond tensors, each consuming one copy leg per endpoint - leg_idx = dict.fromkeys(spins, 1) # leg 0 reserved for field - for (i, j), jij in zip(edges, j_vals): - te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) - ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) - tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) - tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) - leg_idx[i] += 1 - leg_idx[j] += 1 - nodes.append(ten) - return nodes - - -def brute_force_energy(spins, edges, j_vals, h): - """Brute-force min energy: E = -sum J s_i s_j - sum h s_i.""" - n = len(spins) - best = np.inf - for cfg in itertools.product([-1, 1], repeat=n): - e = 0.0 - for (i, j), jij in zip(edges, j_vals): - e -= jij * cfg[i] * cfg[j] - for i, hi in zip(spins, h): - e -= hi * cfg[i] - best = min(best, e) - return best - - -def brute_force_energy_and_degeneracy(n, edges, j_vals, h): - """Brute-force ground energy + degeneracy over spin configs. - - cfg in {0,1}^n; spin s = 1 - 2*cfg (cfg=0 -> s=+1). - E = -sum J s_i s_j - sum h s_i. Returns (best_e, degeneracy). - """ - best_e = None - deg = 0 - for cfg in itertools.product([0, 1], repeat=n): - e = 0.0 - for (i, j), jij in zip(edges, j_vals): - si, sj = 1 - 2 * cfg[i], 1 - 2 * cfg[j] - e -= jij * si * sj - for i, hi in zip(range(n), h): - e -= hi * (1 - 2 * cfg[i]) - if best_e is None or e < best_e - 1e-9: - best_e, deg = e, 1 - elif abs(e - best_e) < 1e-9: - deg += 1 - return best_e, deg - - -# --- Task 10: non-scalar (one free spin) builder + brute force --- - - -# Topology shared by build_ring_with_free_spin / brute_nonscalar_counting. -# A 4-spin ring with spin 0 free: small enough to brute-force exactly, large -# enough to need both tensordot and einsum steps under cotengra. Integer -# couplings/fields keep the brute force exact (no rtol needed). -_RING_SPINS = (0, 1, 2, 3) -_RING_EDGES = ((0, 1), (1, 2), (2, 3), (3, 0)) -_RING_J = (-2, 1, 1, -2) # mixed ferro/antiferro of both magnitudes -> exercises -# # both tensordot and einsum branches under cotengra. -_RING_H = (2, 0, 2, -1) # tuned (via brute-force search) so the two free-spin -# # values give DIFFERENT energy AND DIFFERENT degeneracy: -# # v_free=0 -> E=7, N=2; v_free=1 -> E=5, N=1. A swapped -# # or permuted aux cannot pass the per-output assertion. -_FREE_SPIN = 0 - - -def build_ring_with_free_spin( - spins=_RING_SPINS, - edges=_RING_EDGES, - j_vals=_RING_J, - h=_RING_H, - free_spin=_FREE_SPIN, -): - """Build a small ring Ising with ONE spin's CopyNode given an extra dangling - leg (the free output index). Same ``Tv``/``Te``/``tn.CopyNode`` construction - as ``build_ising_tn`` -- the only difference is the free spin's CopyNode has - one ADDITIONAL leg (beyond degree+1) left unconnected, which becomes the - contraction's output index. - - All legs of a CopyNode share one hyperedge symbol, so leaving the extra leg - dangling means the spin variable s_free appears in the einsum output -- for - each of its two values, cotengra maximizes over the other spins. cfg=0 -> - s=+1, cfg=1 -> s=-1 (matching ``build_ising_tn``'s ``Tv=[h,-h]`` convention). - - Returns ``(nodes, free_edge, j_vals, h)`` where ``free_edge`` is the dangling - CopyNode leg to pass as ``output_edge_order=[free_edge]``. - """ - be = tc.backend - degree = dict.fromkeys(spins, 0) - for i, j in edges: - degree[i] += 1 - degree[j] += 1 - - nodes = [] - copy_nodes = {} - for i in spins: - rank = degree[i] + 1 # one field leg + one per incident edge - if i == free_spin: - rank += 1 # extra leg -> dangling output index for the free spin - cn = tn.CopyNode(rank, 2) - copy_nodes[i] = cn - nodes.append(cn) - # field tensor on leg 0 (same as build_ising_tn) - tv = np.array([h[i], -h[i]], dtype=np.float64) - tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) - tn.connect(cn[0], tvn[0]) - nodes.append(tvn) - - # bond tensors, each consuming one copy leg per endpoint - leg_idx = dict.fromkeys(spins, 1) # leg 0 reserved for field - for (i, j), jij in zip(edges, j_vals): - te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) - ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) - tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) - tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) - leg_idx[i] += 1 - leg_idx[j] += 1 - nodes.append(ten) - - # the free leg is the LAST leg of the free spin's CopyNode - # (legs 0..degree are field+bonds; leg degree+1 is the dangling output). - free_edge = copy_nodes[free_spin][degree[free_spin] + 1] - return nodes, free_edge, j_vals, h - - -def _ising_config_energy(cfg, edges, j_vals, spins, h): - """Ising energy E = -sum J s_i s_j - sum h s_i for a config dict (spin idx -> 0/1).""" - energy = 0.0 - for (i, j), jij in zip(edges, j_vals): - energy -= jij * (1 - 2 * cfg[i]) * (1 - 2 * cfg[j]) - for i, hi in zip(spins, h): - energy -= hi * (1 - 2 * cfg[i]) - return energy - - -def _min_energy_and_degeneracy(energies, eps=1e-9): - """Return ``(min energy, degeneracy)`` -- degeneracy counts configs within eps of min.""" - best = None - deg = 0 - for e in energies: - if best is None or e < best - eps: - best, deg = e, 1 - elif abs(e - best) < eps: - deg += 1 - return best, deg - - -def brute_nonscalar_counting( - nodes, - free_edge, - j_vals, - h, - spins=_RING_SPINS, - edges=_RING_EDGES, - free_spin=_FREE_SPIN, -): - """Brute-force per-output (energy, degeneracy) for the non-scalar ring with - one free spin. Returns ``(expected_e, expected_n)`` -- two arrays of shape - ``(2,)`` indexed by the free spin's value (0 -> s=+1, 1 -> s=-1). - - For each value v of the free spin, enumerate all configs of the OTHER spins, - compute the Ising energy E = -sum J s_i s_j - sum h s_i with s_free fixed, - and record min E + its degeneracy. The contraction returns max-plus(-E) per - output (== -min E), so we return ``-min E`` to match. - - ``nodes`` and ``free_edge`` are accepted only for signature compatibility - with the call site; the brute force is parameterised by the (hardcoded) - ring topology that ``build_ring_with_free_spin`` emits. - """ - del nodes, free_edge # signature-only; topology is fixed by the defaults above - others = [s for s in spins if s != free_spin] - - expected_e = np.zeros(2, dtype=np.float64) - expected_n = np.zeros(2, dtype=np.float64) - for v_free in (0, 1): - fixed = {free_spin: v_free} - energies = [ - _ising_config_energy( - {**fixed, **dict(zip(others, cfg_rest))}, edges, j_vals, spins, h - ) - for cfg_rest in itertools.product([0, 1], repeat=len(others)) - ] - best_e, deg = _min_energy_and_degeneracy(energies) - # max-plus convention: contraction returns max(-E) = -min(E) - expected_e[v_free] = -best_e - expected_n[v_free] = deg - return expected_e, expected_n - - -# --- Task 10 follow-up: non-scalar (two free spins) builder + brute force --- -# -# Extends the one-free-spin canary to a 2D output so the aux-reorder permutation -# in ``cons._algebraic_base_contraction`` (``perm = [dangling_edges.index(e) for e -# in order]``) is exercised with a NON-IDENTITY permutation. With one free edge -# the perm is always ``(0,)`` (identity); with two free edges + a reversed -# ``output_edge_order`` it becomes ``(1, 0)`` (swap), so an elided or -# wrong-direction transpose on the aux tensor is no longer silently a no-op. -# -# Same ring topology / parameters as the one-free-spin case (``_RING_*`` constants -# above), but TWO spins (0 and 1) get an extra dangling CopyNode leg. With -# ``h = [2, 0, 2, -1]`` the per-(a,b) brute-force result is ASYMMETRIC under -# transpose (E = [[5,7],[5,-1]], N = [[1,2],[1,1]]), which is load-bearing: it -# makes ``assert_allclose(N, expected_n_ab.T)`` fail if aux is left in the -# dangling order rather than the ``output_edge_order`` order. -_FREE_SPIN_A = 0 -_FREE_SPIN_B = 1 - - -def build_ring_with_two_free_spins( - spins=_RING_SPINS, - edges=_RING_EDGES, - j_vals=_RING_J, - h=_RING_H, - free_spin_a=_FREE_SPIN_A, - free_spin_b=_FREE_SPIN_B, -): - """Build a small ring Ising with TWO spins' CopyNodes given an extra dangling - leg each (two free output indices). Same ``Tv``/``Te``/``tn.CopyNode`` - construction as ``build_ising_tn`` / ``build_ring_with_free_spin`` -- the only - difference is that BOTH ``free_spin_a`` and ``free_spin_b`` have a CopyNode of - rank ``degree + 2`` (one field leg + one per incident edge + one dangling - output leg). - - The two dangling legs become the contraction's two-axis output. Under - ``cons.sorted_edges`` the dangling order comes out as ``[edge_a, edge_b]`` - (spin 0 before spin 1), so contracting with - ``output_edge_order=[edge_b, edge_a]`` exercises a NON-trivial aux - permutation ``(1, 0)``. - - Returns ``(nodes, edge_a, edge_b, j_vals, h)`` where ``edge_a`` is the - dangling leg of ``free_spin_a`` (spin 0) and ``edge_b`` that of - ``free_spin_b`` (spin 1). Pass them to ``output_edge_order`` in whichever - order the test wants to validate. - """ - be = tc.backend - free_spins = {free_spin_a, free_spin_b} - degree = dict.fromkeys(spins, 0) - for i, j in edges: - degree[i] += 1 - degree[j] += 1 - - nodes = [] - copy_nodes = {} - for i in spins: - rank = degree[i] + 1 # one field leg + one per incident edge - if i in free_spins: - rank += 1 # extra leg -> dangling output index for this free spin - cn = tn.CopyNode(rank, 2) - copy_nodes[i] = cn - nodes.append(cn) - # field tensor on leg 0 (same as build_ising_tn) - tv = np.array([h[i], -h[i]], dtype=np.float64) - tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) - tn.connect(cn[0], tvn[0]) - nodes.append(tvn) - - # bond tensors, each consuming one copy leg per endpoint - leg_idx = dict.fromkeys(spins, 1) # leg 0 reserved for field - for (i, j), jij in zip(edges, j_vals): - te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) - ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) - tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) - tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) - leg_idx[i] += 1 - leg_idx[j] += 1 - nodes.append(ten) - - # each free leg is the LAST leg of its CopyNode - # (legs 0..degree are field+bonds; leg degree+1 is the dangling output). - edge_a = copy_nodes[free_spin_a][degree[free_spin_a] + 1] - edge_b = copy_nodes[free_spin_b][degree[free_spin_b] + 1] - return nodes, edge_a, edge_b, j_vals, h - - -def brute_nonscalar_counting_2d( - nodes, - edge_a, - edge_b, - j_vals, - h, - spins=_RING_SPINS, - edges=_RING_EDGES, - free_spin_a=_FREE_SPIN_A, - free_spin_b=_FREE_SPIN_B, -): - """Brute-force per-output ``(energy, degeneracy)`` for the non-scalar ring - with TWO free spins. Returns ``(expected_e, expected_n)`` -- two ``(2, 2)`` - arrays indexed by ``[v_a, v_b]`` where ``v_a`` is the value of - ``free_spin_a`` (cfg=0 -> s=+1, cfg=1 -> s=-1) and ``v_b`` that of - ``free_spin_b``. - - For each ``(v_a, v_b)`` pair, enumerate the other spins, compute - ``E = -sum J s_i s_j - sum h s_i`` with both free spins fixed, and record - ``min E`` + its degeneracy. Returns ``-min E`` to match the max-plus sign - convention of ``brute_nonscalar_counting``. - - ``nodes`` / ``edge_a`` / ``edge_b`` are signature-only (the topology is fixed - by the defaults above); they are accepted so the call site reads symmetrically - with ``build_ring_with_two_free_spins``'s return signature. - - For the default parameters the brute-force result is - ``E = [[5, 7], [5, -1]]`` and ``N = [[1, 2], [1, 1]]``, both asymmetric under - transpose -- a property the test relies on so a missed aux transpose cannot - pass ``assert_allclose(N, expected_n_ab.T)``. - """ - del nodes, edge_a, edge_b # signature-only; topology is fixed by the defaults above - others = [s for s in spins if s not in (free_spin_a, free_spin_b)] - - expected_e = np.zeros((2, 2), dtype=np.float64) - expected_n = np.zeros((2, 2), dtype=np.float64) - for v_a in (0, 1): - for v_b in (0, 1): - fixed = {free_spin_a: v_a, free_spin_b: v_b} - energies = [ - _ising_config_energy( - {**fixed, **dict(zip(others, cfg_rest))}, edges, j_vals, spins, h - ) - for cfg_rest in itertools.product([0, 1], repeat=len(others)) - ] - best_e, deg = _min_energy_and_degeneracy(energies) - # max-plus convention: contraction returns max(-E) = -min(E) - expected_e[v_a, v_b] = -best_e - expected_n[v_a, v_b] = deg - return expected_e, expected_n diff --git a/tests/test_tropical.py b/tests/test_tropical.py new file mode 100644 index 00000000..be2af826 --- /dev/null +++ b/tests/test_tropical.py @@ -0,0 +1,1139 @@ +"""Tropical (max-plus) and counting tropical algebra tests. + +Covers: max-plus primitives, einsum/tensordot, configuration recovery, +counting (energy + degeneracy), Ising end-to-end, and the non-scalar +aux-reorder canaries. +""" + +import itertools + +import numpy as np +import pytest +import tensornetwork as tn +import opt_einsum + +import tensorcircuit as tc +import tensorcircuit.cons as cons +import tensorcircuit.backends.numpy_backend as nb +import applications.tropical_algebra as tr +from applications.tropical_algebra import ( + MaxPlusAlgebra, + MaxPlusTrackingAlgebra, + CountingTropicalAlgebra, + CountingRepresentation, + tropical, + counting_tropical, + recover_configuration, + get_recorded_topology, + degeneracy, + split_energy_count, + _tropical_tensordot, + _tropical_einsum, + _counting_tensordot, + _counting_einsum, + _td_backtrack_step, +) + +# ═══════════════════════════════════════════════════════════════════════════════ +# Shared utilities (was tests/_tropical_test_utils.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def build_ising_tn(spins, edges, j_vals, h): + """Build the tropical (max-plus) Ising tensor network as tensornetwork nodes. + + Per spin: a ``tn.CopyNode`` of rank (degree+1), dimension 2. A CopyNode is a true + hyperedge hub -- tensornetwork/cotengra identify all its legs as one shared index + (the delta/copy constraint is structural, not stored as data). This is what makes + the contraction route hyperedge-containing pairs through the tropical **einsum** + branch and ordinary pairs through the tropical **tensordot** branch, covering both. + Per edge (i,j): Te[si,sj] = J*si*sj -> [[J,-J],[-J,J]]. + Per spin: field Tv[s] = h*s -> [h,-h]. + Contracting over max-plus yields max_cfg(-E) = -E_ground. + """ + be = tc.backend + degree = dict.fromkeys(spins, 0) + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + + nodes = [] + copy_nodes = {} + for i in spins: + cn = tn.CopyNode(degree[i] + 1, 2) + copy_nodes[i] = cn + nodes.append(cn) + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + leg_idx = dict.fromkeys(spins, 1) + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) + tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) + leg_idx[i] += 1 + leg_idx[j] += 1 + nodes.append(ten) + return nodes + + +def brute_force_energy(spins, edges, j_vals, h): + """Brute-force min energy: E = -sum J s_i s_j - sum h s_i.""" + n = len(spins) + best = np.inf + for cfg in itertools.product([-1, 1], repeat=n): + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + e -= jij * cfg[i] * cfg[j] + for i, hi in zip(spins, h): + e -= hi * cfg[i] + best = min(best, e) + return best + + +def brute_force_energy_and_degeneracy(n, edges, j_vals, h): + """Brute-force ground energy + degeneracy over spin configs. + + cfg in {0,1}^n; spin s = 1 - 2*cfg (cfg=0 -> s=+1). + E = -sum J s_i s_j - sum h s_i. Returns (best_e, degeneracy). + """ + best_e = None + deg = 0 + for cfg in itertools.product([0, 1], repeat=n): + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + si, sj = 1 - 2 * cfg[i], 1 - 2 * cfg[j] + e -= jij * si * sj + for i, hi in zip(range(n), h): + e -= hi * (1 - 2 * cfg[i]) + if best_e is None or e < best_e - 1e-9: + best_e, deg = e, 1 + elif abs(e - best_e) < 1e-9: + deg += 1 + return best_e, deg + + +# --- Non-scalar (free-spin) helpers --- + +_RING_SPINS = (0, 1, 2, 3) +_RING_EDGES = ((0, 1), (1, 2), (2, 3), (3, 0)) +_RING_J = (-2, 1, 1, -2) +_RING_H = (2, 0, 2, -1) +_FREE_SPIN = 0 +_FREE_SPIN_A = 0 +_FREE_SPIN_B = 1 + + +def build_ring_with_free_spin( + spins=_RING_SPINS, + edges=_RING_EDGES, + j_vals=_RING_J, + h=_RING_H, + free_spin=_FREE_SPIN, +): + """Build a small ring Ising with ONE spin's CopyNode given an extra dangling + leg (the free output index).""" + be = tc.backend + degree = dict.fromkeys(spins, 0) + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + + nodes = [] + copy_nodes = {} + for i in spins: + rank = degree[i] + 1 + if i == free_spin: + rank += 1 + cn = tn.CopyNode(rank, 2) + copy_nodes[i] = cn + nodes.append(cn) + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + + leg_idx = dict.fromkeys(spins, 1) + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) + tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) + leg_idx[i] += 1 + leg_idx[j] += 1 + nodes.append(ten) + + free_edge = copy_nodes[free_spin][degree[free_spin] + 1] + return nodes, free_edge, j_vals, h + + +def build_ring_with_two_free_spins( + spins=_RING_SPINS, + edges=_RING_EDGES, + j_vals=_RING_J, + h=_RING_H, + free_spin_a=_FREE_SPIN_A, + free_spin_b=_FREE_SPIN_B, +): + """Build a small ring Ising with TWO spins' CopyNodes given an extra dangling + leg each (two free output indices).""" + be = tc.backend + free_spins = {free_spin_a, free_spin_b} + degree = dict.fromkeys(spins, 0) + for i, j in edges: + degree[i] += 1 + degree[j] += 1 + + nodes = [] + copy_nodes = {} + for i in spins: + rank = degree[i] + 1 + if i in free_spins: + rank += 1 + cn = tn.CopyNode(rank, 2) + copy_nodes[i] = cn + nodes.append(cn) + tv = np.array([h[i], -h[i]], dtype=np.float64) + tvn = tn.Node(be.cast(be.convert_to_tensor(tv), "float64")) + tn.connect(cn[0], tvn[0]) + nodes.append(tvn) + + leg_idx = dict.fromkeys(spins, 1) + for (i, j), jij in zip(edges, j_vals): + te = np.array([[jij, -jij], [-jij, jij]], dtype=np.float64) + ten = tn.Node(be.cast(be.convert_to_tensor(te), "float64")) + tn.connect(ten[0], copy_nodes[i][leg_idx[i]]) + tn.connect(ten[1], copy_nodes[j][leg_idx[j]]) + leg_idx[i] += 1 + leg_idx[j] += 1 + nodes.append(ten) + + edge_a = copy_nodes[free_spin_a][degree[free_spin_a] + 1] + edge_b = copy_nodes[free_spin_b][degree[free_spin_b] + 1] + return nodes, edge_a, edge_b, j_vals, h + + +def _ising_config_energy(cfg, edges, j_vals, spins, h): + """Ising energy E = -sum J s_i s_j - sum h s_i for a config dict (spin idx -> 0/1).""" + energy = 0.0 + for (i, j), jij in zip(edges, j_vals): + energy -= jij * (1 - 2 * cfg[i]) * (1 - 2 * cfg[j]) + for i, hi in zip(spins, h): + energy -= hi * (1 - 2 * cfg[i]) + return energy + + +def _min_energy_and_degeneracy(energies, eps=1e-9): + """Return ``(min energy, degeneracy)`` -- degeneracy counts configs within eps of min.""" + best = None + deg = 0 + for e in energies: + if best is None or e < best - eps: + best, deg = e, 1 + elif abs(e - best) < eps: + deg += 1 + return best, deg + + +def brute_nonscalar_counting( + nodes, + free_edge, + j_vals, + h, + spins=_RING_SPINS, + edges=_RING_EDGES, + free_spin=_FREE_SPIN, +): + """Brute-force per-output (energy, degeneracy) for the non-scalar ring with + one free spin.""" + del nodes, free_edge + others = [s for s in spins if s != free_spin] + + expected_e = np.zeros(2, dtype=np.float64) + expected_n = np.zeros(2, dtype=np.float64) + for v_free in (0, 1): + fixed = {free_spin: v_free} + energies = [ + _ising_config_energy( + {**fixed, **dict(zip(others, cfg_rest))}, edges, j_vals, spins, h + ) + for cfg_rest in itertools.product([0, 1], repeat=len(others)) + ] + best_e, deg = _min_energy_and_degeneracy(energies) + expected_e[v_free] = -best_e + expected_n[v_free] = deg + return expected_e, expected_n + + +def brute_nonscalar_counting_2d( + nodes, + edge_a, + edge_b, + j_vals, + h, + spins=_RING_SPINS, + edges=_RING_EDGES, + free_spin_a=_FREE_SPIN_A, + free_spin_b=_FREE_SPIN_B, +): + """Brute-force per-output (energy, degeneracy) for the non-scalar ring with + TWO free spins.""" + del nodes, edge_a, edge_b + others = [s for s in spins if s not in (free_spin_a, free_spin_b)] + + expected_e = np.zeros((2, 2), dtype=np.float64) + expected_n = np.zeros((2, 2), dtype=np.float64) + for v_a in (0, 1): + for v_b in (0, 1): + fixed = {free_spin_a: v_a, free_spin_b: v_b} + energies = [ + _ising_config_energy( + {**fixed, **dict(zip(others, cfg_rest))}, edges, j_vals, spins, h + ) + for cfg_rest in itertools.product([0, 1], repeat=len(others)) + ] + best_e, deg = _min_energy_and_degeneracy(energies) + expected_e[v_a, v_b] = -best_e + expected_n[v_a, v_b] = deg + return expected_e, expected_n + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Section 1 — Core algebra: tensordot + einsum (was test_tropical_algebra.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def _be(): + return tc.backend + + +def _ref_tropical_tensordot(anp, bnp, axes): + if isinstance(axes, int): + a_axes = list(range(anp.ndim - axes, anp.ndim)) + b_axes = list(range(0, axes)) + else: + a_axes, b_axes = list(axes[0]), list(axes[1]) + a_free = [i for i in range(anp.ndim) if i not in a_axes] + b_free = [i for i in range(bnp.ndim) if i not in b_axes] + a_t = np.transpose(anp, a_free + a_axes) + b_t = np.transpose(bnp, b_axes + b_free) + a_fs = [anp.shape[i] for i in a_free] + b_fs = [bnp.shape[i] for i in b_free] + a2 = a_t.reshape(-1, np.prod([anp.shape[i] for i in a_axes], dtype=int) or 1) + b2 = b_t.reshape(np.prod([bnp.shape[i] for i in b_axes], dtype=int) or 1, -1) + A, b_n = a2.shape[0], b2.shape[1] + res = np.full((A, b_n), -np.inf) + for i in range(A): + for j in range(b_n): + res[i, j] = np.max(a2[i, :] + b2[:, j]) + return res.reshape(tuple(a_fs) + tuple(b_fs)) + + +def _ref_tropical_einsum(eq, a, b): + lhs, rhs = eq.split("->") + ia, ib = lhs.split(",") + sizes = {} + for s, t in zip([ia, ib], [a, b]): + for c, dim in zip(s, t.shape): + sizes[c] = dim + out = np.full([sizes[c] for c in rhs], -np.inf) + allc = list(dict.fromkeys(list(ia) + list(ib))) + for combo in itertools.product(*[range(sizes[c]) for c in allc]): + env = dict(zip(allc, combo)) + ia_idx = tuple(env[c] for c in ia) + ib_idx = tuple(env[c] for c in ib) + val = a[ia_idx] + b[ib_idx] + oidx = tuple(env[c] for c in rhs) + if val > out[oidx]: + out[oidx] = val + return out + + +# --- tensordot tests --- + + +def test_tropical_tensordot_matrix(): + be = _be() + rng = np.random.default_rng(0) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_tensordot(be, a, b, axes=1)) + np.testing.assert_allclose(got, _ref_tropical_tensordot(anp, bnp, 1)) + + +def test_tropical_tensordot_multi_axis(): + be = _be() + rng = np.random.default_rng(1) + anp, bnp = rng.normal(size=(2, 3, 4)), rng.normal(size=(3, 4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_tensordot(be, a, b, axes=([1, 2], [0, 1]))) + np.testing.assert_allclose(got, _ref_tropical_tensordot(anp, bnp, ([1, 2], [0, 1]))) + assert got.shape == (2, 5) + + +# --- einsum tests --- + + +def test_tropical_einsum_pair(): + be = _be() + rng = np.random.default_rng(2) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "ab,bc->ac", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("ab,bc->ac", anp, bnp)) + + +def test_tropical_einsum_hyperedge_shared_index(): + be = _be() + rng = np.random.default_rng(3) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "ab,ac->abc", a, b)) + ref = _ref_tropical_einsum("ab,ac->abc", anp, bnp) + np.testing.assert_allclose(got, ref) + + +def test_maxplus_algebra_uses_tropical_ops(): + be = _be() + rng = np.random.default_rng(4) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + alg = MaxPlusAlgebra() + assert alg.name == "maxplus" + np.testing.assert_allclose( + np.array(alg.tensordot(be, a, b, 1)), + _ref_tropical_tensordot(anp, bnp, 1), + ) + np.testing.assert_allclose( + np.array(alg.einsum(be, "ab,bc->ac", a, b)), + _ref_tropical_einsum("ab,bc->ac", anp, bnp), + ) + + +def test_tropical_einsum_single_tensor_transpose_ok(): + be = _be() + rng = np.random.default_rng(5) + anp = rng.normal(size=(3, 4)) + a = be.cast(be.convert_to_tensor(anp), "float64") + got = np.array(_tropical_einsum(be, "ab->ba", a)) + np.testing.assert_allclose(got, np.array(be.einsum("ab->ba", a))) + + +def test_tropical_einsum_single_tensor_trace(): + be = _be() + rng = np.random.default_rng(6) + square_a = rng.normal(size=(3, 3)) + a = be.cast(be.convert_to_tensor(square_a), "float64") + got = float(np.array(_tropical_einsum(be, "ii->", a))) + assert np.isclose(got, np.max(np.diag(square_a))) + + +def test_tropical_einsum_single_tensor_diagonal(): + be = _be() + rng = np.random.default_rng(7) + square_a = rng.normal(size=(4, 4)) + a = be.cast(be.convert_to_tensor(square_a), "float64") + got = np.array(_tropical_einsum(be, "ii->i", a)) + np.testing.assert_allclose(got, np.diag(square_a)) + + +def test_tropical_einsum_single_tensor_reduce(): + be = _be() + rng = np.random.default_rng(8) + anp = rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + got = np.array(_tropical_einsum(be, "ab->a", a)) + np.testing.assert_allclose(got, np.max(anp, axis=1)) + + +def test_tropical_einsum_single_tensor_partial_trace(): + be = _be() + rng = np.random.default_rng(9) + anp = rng.normal(size=(4, 4, 3)) + a = be.cast(be.convert_to_tensor(anp), "float64") + got = np.array(_tropical_einsum(be, "iij->j", a)) + ref = np.max(np.array([anp[i, i, :] for i in range(4)]), axis=0) + np.testing.assert_allclose(got, ref) + + +def test_tropical_einsum_intra_operand_repeat_first(): + be = _be() + rng = np.random.default_rng(20) + anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "iij,jk->ik", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jk->ik", anp, bnp)) + + +def test_tropical_einsum_intra_operand_repeat_contracted(): + be = _be() + rng = np.random.default_rng(21) + anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "iij,jk->k", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jk->k", anp, bnp)) + + +def test_tropical_einsum_both_operands_repeat(): + be = _be() + rng = np.random.default_rng(22) + anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 4)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(_tropical_einsum(be, "iij,jj->ij", a, b)) + np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jj->ij", anp, bnp)) + + +def test_tropical_context_changes_result(): + be = _be() + anp = np.array([[1.0, 5.0], [3.0, 2.0]]) + bnp = np.array([10.0, 0.0]) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + na, nb = tn.Node(a), tn.Node(b) + tn.connect(na[1], nb[0]) + expected = np.array([max(1 + 10, 5 + 0), max(3 + 10, 2 + 0)]) + with tropical(): + got = np.array(cons.contractor([na, nb], output_edge_order=[na[0]]).tensor) + np.testing.assert_allclose(got, expected) + + +def test_preprocessing_does_not_corrupt_tropical(): + be = _be() + rng = np.random.default_rng(0) + n0 = rng.normal(size=(2,)) + n1 = rng.normal(size=(2, 2)) + n2 = rng.normal(size=(2, 2)) + n3 = rng.normal(size=(2, 2)) + n4 = rng.normal(size=(2,)) + na = tn.Node(be.cast(be.convert_to_tensor(n0), "float64")) + nb = tn.Node(be.cast(be.convert_to_tensor(n1), "float64")) + nc = tn.Node(be.cast(be.convert_to_tensor(n2), "float64")) + nd = tn.Node(be.cast(be.convert_to_tensor(n3), "float64")) + ne = tn.Node(be.cast(be.convert_to_tensor(n4), "float64")) + tn.connect(na[0], nb[0]) + tn.connect(nb[1], nc[0]) + tn.connect(nc[1], nd[0]) + tn.connect(nd[1], ne[0]) + nodes = [na, nb, nc, nd, ne] + assert len(nodes) >= 5 + + ref = -np.inf + for s0, s1, s2, s3 in itertools.product(range(2), repeat=4): + val = n0[s0] + n1[s0, s1] + n2[s1, s2] + n3[s2, s3] + n4[s3] + ref = max(ref, val) + + with tropical(): + got = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + + np.testing.assert_allclose(got, ref, atol=1e-6) + + +def test_tropical_public_api_surface(): + for name in [ + "MaxPlusAlgebra", + "MaxPlusTrackingAlgebra", + "CountingTropicalAlgebra", + "tropical", + "counting_tropical", + "recover_configuration", + "split_energy_count", + ]: + assert hasattr(tr, name), name + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Section 2 — Configuration recovery (was test_tropical_config.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def _brute_cfg(n, edges, j_vals, h): + """Brute-force min energy over spin configs; returns (best_e, set_of_optimal_cfgs).""" + best_e = None + best_cfgs = set() + for cfg in itertools.product([0, 1], repeat=n): + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + e -= jij * (1 - 2 * cfg[i]) * (1 - 2 * cfg[j]) + for i, hi in zip(range(n), h): + e -= hi * (1 - 2 * cfg[i]) + if best_e is None or e < best_e - 1e-9: + best_e = e + best_cfgs = {cfg} + elif abs(e - best_e) < 1e-9: + best_cfgs.add(cfg) + return best_e, best_cfgs + + +def _energy_of(assignment, input_sets, raw_tensors): + """Max-plus total for a symbol->value assignment = -E for that config.""" + total = 0.0 + for term, t in zip(input_sets, raw_tensors): + idx = tuple(int(assignment[c]) for c in term) + total += float(np.asarray(t)[idx]) + return total + + +def _contract_tropical_track(nodes): + """Contract under tracking tropical; return (value, config).""" + with tropical(track=True): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + cfg = recover_configuration() + return val, cfg + + +def test_config_recovery_tiny_ising_unique(): + spins = [0, 1] + edges = [(0, 1)] + J = [0.7] + h = [0.3, -0.2] + best_e, best_cfgs = _brute_cfg(len(spins), edges, J, h) + + nodes = build_ising_tn(spins, edges, J, h) + val, cfg = _contract_tropical_track(nodes) + + np.testing.assert_allclose(val, -best_e, atol=1e-6) + + _tree, input_sets, raw_tensors = get_recorded_topology() + assert input_sets is not None and raw_tensors is not None + all_labels = set().union(*[set(t) for t in input_sets]) + assert set(cfg.keys()) == all_labels + + recovered_total = _energy_of(cfg, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, val, atol=1e-6) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg[term[0]]) + assert tuple(spin_cfg[i] for i in range(len(spins))) in best_cfgs + + +def test_config_recovery_tiny_ising_degenerate(): + spins = [0, 1] + edges = [(0, 1)] + J = [1.0] + h = [0.0, 0.0] + best_e, best_cfgs = _brute_cfg(len(spins), edges, J, h) + assert len(best_cfgs) == 2 + + nodes = build_ising_tn(spins, edges, J, h) + val, cfg = _contract_tropical_track(nodes) + + np.testing.assert_allclose(val, -best_e, atol=1e-6) + _tree, input_sets, raw_tensors = get_recorded_topology() + recovered_total = _energy_of(cfg, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg[term[0]]) + assert tuple(spin_cfg[i] for i in range(len(spins))) in best_cfgs + + +def test_config_recovery_three_ring(): + spins = [0, 1, 2] + edges = [(0, 1), (1, 2)] + J = [0.5, -0.8] + h = [0.1, 0.2, -0.3] + best_e, _best_cfgs = _brute_cfg(len(spins), edges, J, h) + + nodes = build_ising_tn(spins, edges, J, h) + val, cfg = _contract_tropical_track(nodes) + np.testing.assert_allclose(val, -best_e, atol=1e-6) + + _tree, input_sets, raw_tensors = get_recorded_topology() + recovered_total = _energy_of(cfg, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + +def test_config_recovery_non_scalar_raises(): + be = tc.backend + a_np = np.array([[1.0, 4.0, 2.0], [3.0, 1.0, 5.0], [0.0, 2.0, 1.0]]) + b_np = np.array([0.5, 1.0, -0.5]) + + a_node = tn.Node(be.cast(be.convert_to_tensor(a_np), "float64")) + b_node = tn.Node(be.cast(be.convert_to_tensor(b_np), "float64")) + tn.connect(a_node[1], b_node[0]) + nodes = [a_node, b_node] + + with tropical(track=True): + res = cons.contractor(nodes, output_edge_order=[a_node[0]]) + assert np.array(res.tensor).shape == (3,) + with pytest.raises(NotImplementedError): + recover_configuration() + + +class _StubTree: + """Minimal stand-in for a cotengra ``ContractionTree``.""" + + def __init__(self, inds, axes, perm): + self._inds = inds + self._axes = axes + self._perm = perm + + def get_inds(self, node): + return self._inds[node] + + def get_tensordot_axes(self, node): + return self._axes[node] + + def get_tensordot_perm(self, node): + return self._perm[node] + + +def test_td_backtrack_step_perm_inversion(): + l_node, r_node, p_node = "l", "r", "p" + tree = _StubTree( + inds={l_node: "ab", r_node: "bcd", p_node: "dac"}, + axes={p_node: ([1], [0])}, + perm={p_node: (2, 0, 1)}, + ) + argmax = np.zeros((2, 3, 2), dtype=np.int64) + argmax[1, 2, 0] = 3 + rec = {"kind": "td", "argmax": argmax, "contract_dims": (4,)} + p_pos = (0, 1, 2) + assignment = {} + l_pos, r_pos = _td_backtrack_step( + tree, p_node, l_node, r_node, rec, p_pos, assignment + ) + + assert assignment["b"] == 3 + assert l_pos == (1, 3) + assert r_pos == (3, 2, 0) + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Section 3 — Counting (was test_tropical_counting.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def _ref_counting_tensordot(anp, bnp, axes, eps=1e-9): + if isinstance(axes, int): + a_axes = list(range(anp.ndim - axes, anp.ndim)) + b_axes = list(range(0, axes)) + else: + a_axes, b_axes = list(axes[0]), list(axes[1]) + a_free = [i for i in range(anp.ndim) if i not in a_axes] + b_free = [i for i in range(bnp.ndim) if i not in b_axes] + at = np.transpose(anp, a_free + a_axes) + bt = np.transpose(bnp, b_axes + b_free) + a_fs = [anp.shape[i] for i in a_free] + b_fs = [bnp.shape[i] for i in b_free] + k = int(np.prod([anp.shape[i] for i in a_axes]) or 1) + A = int(np.prod(a_fs) or 1) + b_n = int(np.prod(b_fs) or 1) + a2 = at.reshape(A, k) + b2 = bt.reshape(k, b_n) + out_e = np.full((A, b_n), -np.inf) + out_n = np.zeros((A, b_n)) + for i in range(A): + for j in range(b_n): + s = a2[i, :] + b2[:, j] + mx = np.max(s) + out_e[i, j] = mx + out_n[i, j] = int(np.sum(np.abs(s - mx) < eps)) + return out_e.reshape(a_fs + b_fs), out_n.reshape(a_fs + b_fs) + + +def _ref_counting_einsum(eq, a, b, eps=1e-9): + lhs, rhs = eq.split("->") + ia, ib = lhs.split(",") + sizes = {} + for s, t in zip([ia, ib], [a, b]): + for c, dim in zip(s, t.shape): + sizes[c] = dim + out_e = np.full([sizes[c] for c in rhs], -np.inf) + out_n = np.zeros([sizes[c] for c in rhs]) + allc = list(dict.fromkeys(list(ia) + list(ib))) + for combo in itertools.product(*[range(sizes[c]) for c in allc]): + env = dict(zip(allc, combo)) + e = a[tuple(env[c] for c in ia)] + b[tuple(env[c] for c in ib)] + oidx = tuple(env[c] for c in rhs) + if e > out_e[oidx] + eps: + out_e[oidx] = e + out_n[oidx] = 1 + elif abs(e - out_e[oidx]) < eps: + out_n[oidx] += 1 + return out_e, out_n + + +def test_counting_tensordot_matrix(): + be = _be() + rng = np.random.default_rng(10) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast( + be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], axis=-1)), "float64" + ) + b = be.cast( + be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], axis=-1)), "float64" + ) + got = np.array(_counting_tensordot(be, a, b, axes=1)) + ref_e, ref_n = _ref_counting_tensordot(anp, bnp, 1) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_split_energy_count(): + be = _be() + arr = np.array([[1.0, 2.0], [3.0, 4.0]]) + stacked = be.cast( + be.convert_to_tensor(np.stack([arr, arr * 2], axis=-1)), "float64" + ) + e, n = split_energy_count(stacked) + np.testing.assert_allclose(np.array(e), arr) + np.testing.assert_allclose(np.array(n), arr * 2) + + +def test_counting_einsum_pair(): + be = _be() + rng = np.random.default_rng(11) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,bc->ac", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,bc->ac", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_counting_einsum_hyperedge(): + be = _be() + rng = np.random.default_rng(12) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,ac->abc", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,ac->abc", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_counting_einsum_tie_degeneracy(): + be = _be() + anp = np.array([[1.0, 1.0]]) + bnp = np.array([[2.0], [2.0]]) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,bc->ac", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,bc->ac", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + assert ref_n[0, 0] == 2 + np.testing.assert_allclose(got[0, 0, 0], 3.0) + np.testing.assert_allclose(got[0, 0, 1], 2.0) + + +def test_counting_einsum_single_operand_diagonal(): + be = nb.NumpyBackend() + e = np.array([[1.0, 5.0], [3.0, 2.0]]) + n = np.array([[1.0, 2.0], [1.0, 1.0]]) + a = np.stack([e, n], axis=-1) + out = _counting_einsum(be, "aa->a", a) + np.testing.assert_allclose(out[..., 0], np.diag(e)) + np.testing.assert_allclose(out[..., 1], np.diag(n)) + + +def test_counting_einsum_single_operand_reduce(): + be = nb.NumpyBackend() + e = np.array([[1.0, 5.0], [3.0, 3.0]]) + n = np.array([[1.0, 2.0], [1.0, 4.0]]) + a = np.stack([e, n], axis=-1) + out = _counting_einsum(be, "ab->a", a) + np.testing.assert_allclose(out[..., 0], [5.0, 3.0]) + np.testing.assert_allclose(out[..., 1], [2.0, 5.0]) + + +def test_counting_einsum_two_operand_intra_repeat(): + be = nb.NumpyBackend() + e_a = np.array([[2.0, 1.0], [1.0, 2.0]]) + n_a = np.ones((2, 2)) + a = np.stack([e_a, n_a], axis=-1) + e_b = np.array([0.0, 0.0]) + n_b = np.array([1.0, 1.0]) + b = np.stack([e_b, n_b], axis=-1) + out = _counting_einsum(be, "aa,b->ab", a, b) + np.testing.assert_allclose(out[..., 0], [[2, 2], [2, 2]]) + np.testing.assert_allclose(out[..., 1], [[1, 1], [1, 1]]) + + +def test_counting_einsum_multi_axis(): + be = _be() + rng = np.random.default_rng(14) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) + a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") + b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") + got = np.array(_counting_einsum(be, "ab,ac->c", a, b)) + ref_e, ref_n = _ref_counting_einsum("ab,ac->c", anp, bnp) + np.testing.assert_allclose(got[..., 0], ref_e) + np.testing.assert_allclose(got[..., 1], ref_n) + + +def test_degeneracy_none_after_standard_contraction(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, 1.0, 1.0, 1.0, 1.0] + h = [0.0, 0.0, 0.0, 0.0, 0.0] + expected_e, expected_n = brute_force_energy_and_degeneracy(len(spins), edges, J, h) + assert expected_n == 2 + nodes = build_ising_tn(spins, edges, J, h) + + with counting_tropical(): + node = cons._algebraic_base_contraction( + nodes, + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[], + ignore_edge_order=True, + ) + np.testing.assert_allclose(np.array(degeneracy()), expected_n) + np.testing.assert_allclose(np.array(node.tensor), -expected_e, atol=1e-6) + + rng = np.random.default_rng(0) + a = tn.Node(rng.standard_normal((2, 3)).astype(np.complex64)) + b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) + tn.connect(a[1], b[0]) + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], b[1]], + ) + assert degeneracy() is None + + +def test_nonscalar_counting_energy_and_degeneracy_per_output(): + nodes, free_edge, J, h = build_ring_with_free_spin() + expected_e, expected_n = brute_nonscalar_counting(nodes, free_edge, J, h) + + with counting_tropical(): + node = cons._algebraic_base_contraction( + nodes, + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[free_edge], + ) + E = np.asarray(node.tensor) + N = np.asarray(degeneracy()) + + np.testing.assert_allclose(E, expected_e) + np.testing.assert_allclose(N, expected_n) + assert E[0] != E[1] or N[0] != N[1] + + +def test_nonscalar_counting_two_free_spins_reversed_order(): + nodes, edge_a, edge_b, J, h = build_ring_with_two_free_spins() + expected_e_ab, expected_n_ab = brute_nonscalar_counting_2d( + nodes, edge_a, edge_b, J, h + ) + assert not np.allclose( + expected_e_ab, expected_e_ab.T + ), "weak canary: brute-force energy is symmetric under transpose" + assert not np.allclose( + expected_n_ab, expected_n_ab.T + ), "weak canary: brute-force degeneracy is symmetric under transpose" + + with counting_tropical(): + node = cons._algebraic_base_contraction( + nodes, + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[edge_b, edge_a], + ) + E = np.asarray(node.tensor) + N = np.asarray(degeneracy()) + + np.testing.assert_allclose(E, expected_e_ab.T) + np.testing.assert_allclose(N, expected_n_ab.T) + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Section 4 — Example-level tests (was test_tropical_example.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def test_example_maxplus_tensordot_matches_brute(): + be = tc.backend + rng = np.random.default_rng(0) + anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) + a = be.cast(be.convert_to_tensor(anp), "float64") + b = be.cast(be.convert_to_tensor(bnp), "float64") + got = np.array(MaxPlusAlgebra().tensordot(be, a, b, 1)) + ref = np.max(anp[:, :, None] + bnp[None, :, :], axis=1) + np.testing.assert_allclose(got, ref) + + +def test_example_recover_configuration_on_tiny_ising(): + be = tc.backend + n = tn.Node(be.cast(be.convert_to_tensor(np.array([0.0, 0.0])), "float64")) + e = tn.Node( + be.cast(be.convert_to_tensor(np.array([[1.0, -1.0], [-1.0, 1.0]])), "float64") + ) + m = tn.Node(be.cast(be.convert_to_tensor(np.array([0.5, -0.5])), "float64")) + tn.connect(n[0], e[0]) + tn.connect(e[1], m[0]) + nodes = [n, e, m] + with tropical(track=True): + float(cons.contractor(nodes, output_edge_order=[]).tensor) + cfg = recover_configuration() + assert isinstance(cfg, dict) + assert len(cfg) >= 1 + + +# ═══════════════════════════════════════════════════════════════════════════════ +# Section 5 — Ising end-to-end (was test_tropical_ising.py) +# ═══════════════════════════════════════════════════════════════════════════════ + + +def _ising_energy_of_cfg(cfg, edges, j_vals, h): + """Ising energy of a {0,1}-config: E = -sum J s_i s_j - sum h s_i (s=1-2cfg).""" + e = 0.0 + for (i, j), jij in zip(edges, j_vals): + si, sj = 1 - 2 * cfg[i], 1 - 2 * cfg[j] + e -= jij * si * sj + for i, hi in zip(range(len(h)), h): + e -= hi * (1 - 2 * cfg[i]) + return e + + +def _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors): + """Max-plus total of a symbol->value assignment = sum of tensor entries = + -E for that configuration.""" + total = 0.0 + for term, t in zip(input_sets, raw_tensors): + idx = tuple(int(cfg_sym[c]) for c in term) + total += float(np.asarray(t)[idx]) + return total + + +def test_ising_ring_ground_state_matches_bruteforce(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + nodes = build_ising_tn(spins, edges, J, h) + + e_ground = brute_force_energy(spins, edges, J, h) + + with tropical(): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + + np.testing.assert_allclose(val, -e_ground, atol=1e-6) + + +def test_ising_config_recovery_five_ring(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + + best_e = None + for cfg in itertools.product([0, 1], repeat=len(spins)): + e = _ising_energy_of_cfg(cfg, edges, J, h) + if best_e is None or e < best_e - 1e-9: + best_e = e + + nodes = build_ising_tn(spins, edges, J, h) + with tropical(track=True): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + cfg_sym = recover_configuration() + + np.testing.assert_allclose(val, -best_e, atol=1e-6) + + _tree, input_sets, raw_tensors = get_recorded_topology() + assert input_sets is not None and raw_tensors is not None + all_labels = set().union(*[set(t) for t in input_sets]) + assert set(cfg_sym.keys()) == all_labels + + recovered_total = _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, val, atol=1e-6) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg_sym[term[0]]) + assert len(spin_cfg) == len(spins) + recovered_e = _ising_energy_of_cfg( + [spin_cfg[i] for i in range(len(spins))], edges, J, h + ) + np.testing.assert_allclose(recovered_e, best_e, atol=1e-6) + + +def test_ising_config_recovery_degenerate_ring(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, 1.0, 1.0, 1.0, 1.0] + h = [0.0, 0.0, 0.0, 0.0, 0.0] + + best_e = None + for cfg in itertools.product([0, 1], repeat=len(spins)): + e = _ising_energy_of_cfg(cfg, edges, J, h) + if best_e is None or e < best_e - 1e-9: + best_e = e + assert best_e == -5.0 + + nodes = build_ising_tn(spins, edges, J, h) + with tropical(track=True): + val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() + cfg_sym = recover_configuration() + + np.testing.assert_allclose(val, -best_e, atol=1e-6) + _tree, input_sets, raw_tensors = get_recorded_topology() + recovered_total = _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors) + np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) + + spin_cfg = {} + for term, _t in zip(input_sets, raw_tensors): + if len(term) == 1: + spin_cfg[len(spin_cfg)] = int(cfg_sym[term[0]]) + recovered_e = _ising_energy_of_cfg( + [spin_cfg[i] for i in range(len(spins))], edges, J, h + ) + np.testing.assert_allclose(recovered_e, best_e, atol=1e-6) + + +def test_ising_counting_ground_state_and_degeneracy(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, -1.0, 1.0, 1.0, -1.0] + h = [0.5, -0.3, 0.2, 0.0, 0.4] + best_e, deg = brute_force_energy_and_degeneracy(len(spins), edges, J, h) + + nodes = build_ising_tn(spins, edges, J, h) + + with counting_tropical(): + res = cons.contractor(nodes, output_edge_order=[], ignore_edge_order=True) + count = degeneracy() + + energy = np.array(res.tensor) + np.testing.assert_allclose(energy, -best_e, atol=1e-6) + np.testing.assert_allclose(np.array(count), deg, atol=1e-6) + + +def test_ising_counting_degenerate_ring(): + spins = [0, 1, 2, 3, 4] + edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] + J = [1.0, 1.0, 1.0, 1.0, 1.0] + h = [0.0, 0.0, 0.0, 0.0, 0.0] + best_e, deg = brute_force_energy_and_degeneracy(len(spins), edges, J, h) + + assert best_e == -5.0 + assert deg == 2 + + nodes = build_ising_tn(spins, edges, J, h) + + with counting_tropical(): + res = cons.contractor(nodes, output_edge_order=[], ignore_edge_order=True) + count = degeneracy() + + energy = np.array(res.tensor) + np.testing.assert_allclose(energy, -best_e, atol=1e-6) + np.testing.assert_allclose(np.array(count), deg, atol=1e-6) diff --git a/tests/test_tropical_algebra.py b/tests/test_tropical_algebra.py deleted file mode 100644 index 05ab7da0..00000000 --- a/tests/test_tropical_algebra.py +++ /dev/null @@ -1,272 +0,0 @@ -import numpy as np -import tensornetwork as tn -import tensorcircuit as tc -import tensorcircuit.cons as cons -from applications.tropical_algebra import _tropical_tensordot - - -def _be(): - return tc.backend - - -def _ref_tropical_tensordot(anp, bnp, axes): - if isinstance(axes, int): - a_axes = list(range(anp.ndim - axes, anp.ndim)) - b_axes = list(range(0, axes)) - else: - a_axes, b_axes = list(axes[0]), list(axes[1]) - a_free = [i for i in range(anp.ndim) if i not in a_axes] - b_free = [i for i in range(bnp.ndim) if i not in b_axes] - a_t = np.transpose(anp, a_free + a_axes) - b_t = np.transpose(bnp, b_axes + b_free) - a_fs = [anp.shape[i] for i in a_free] - b_fs = [bnp.shape[i] for i in b_free] - a2 = a_t.reshape(-1, np.prod([anp.shape[i] for i in a_axes], dtype=int) or 1) - b2 = b_t.reshape(np.prod([bnp.shape[i] for i in b_axes], dtype=int) or 1, -1) - A, b_n = a2.shape[0], b2.shape[1] - res = np.full((A, b_n), -np.inf) - for i in range(A): - for j in range(b_n): - res[i, j] = np.max(a2[i, :] + b2[:, j]) - return res.reshape(tuple(a_fs) + tuple(b_fs)) - - -def test_tropical_tensordot_matrix(): - be = _be() - rng = np.random.default_rng(0) - anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) - a = be.cast(be.convert_to_tensor(anp), "float64") - b = be.cast(be.convert_to_tensor(bnp), "float64") - got = np.array(_tropical_tensordot(be, a, b, axes=1)) - np.testing.assert_allclose(got, _ref_tropical_tensordot(anp, bnp, 1)) - - -def test_tropical_tensordot_multi_axis(): - be = _be() - rng = np.random.default_rng(1) - anp, bnp = rng.normal(size=(2, 3, 4)), rng.normal(size=(3, 4, 5)) - a = be.cast(be.convert_to_tensor(anp), "float64") - b = be.cast(be.convert_to_tensor(bnp), "float64") - got = np.array(_tropical_tensordot(be, a, b, axes=([1, 2], [0, 1]))) - np.testing.assert_allclose(got, _ref_tropical_tensordot(anp, bnp, ([1, 2], [0, 1]))) - assert got.shape == (2, 5) - - -from applications.tropical_algebra import _tropical_einsum, MaxPlusAlgebra, tropical - - -def _ref_tropical_einsum(eq, a, b): - import itertools - - lhs, rhs = eq.split("->") - ia, ib = lhs.split(",") - sizes = {} - for s, t in zip([ia, ib], [a, b]): - for c, dim in zip(s, t.shape): - sizes[c] = dim - out = np.full([sizes[c] for c in rhs], -np.inf) - allc = list(dict.fromkeys(list(ia) + list(ib))) - for combo in itertools.product(*[range(sizes[c]) for c in allc]): - env = dict(zip(allc, combo)) - ia_idx = tuple(env[c] for c in ia) - ib_idx = tuple(env[c] for c in ib) - val = a[ia_idx] + b[ib_idx] - oidx = tuple(env[c] for c in rhs) - if val > out[oidx]: - out[oidx] = val - return out - - -def test_tropical_einsum_pair(): - be = _be() - rng = np.random.default_rng(2) - anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) - a = be.cast(be.convert_to_tensor(anp), "float64") - b = be.cast(be.convert_to_tensor(bnp), "float64") - got = np.array(_tropical_einsum(be, "ab,bc->ac", a, b)) - np.testing.assert_allclose(got, _ref_tropical_einsum("ab,bc->ac", anp, bnp)) - - -def test_tropical_einsum_hyperedge_shared_index(): - # 'a' shared (batch, kept in output) — this is the copy-node pair shape - be = _be() - rng = np.random.default_rng(3) - anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) - a = be.cast(be.convert_to_tensor(anp), "float64") - b = be.cast(be.convert_to_tensor(bnp), "float64") - got = np.array(_tropical_einsum(be, "ab,ac->abc", a, b)) # note: output order a,b,c - ref = _ref_tropical_einsum("ab,ac->abc", anp, bnp) - np.testing.assert_allclose(got, ref) - - -def test_maxplus_algebra_uses_tropical_ops(): - be = _be() - rng = np.random.default_rng(4) - anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) - a = be.cast(be.convert_to_tensor(anp), "float64") - b = be.cast(be.convert_to_tensor(bnp), "float64") - alg = MaxPlusAlgebra() - assert alg.name == "maxplus" - np.testing.assert_allclose( - np.array(alg.tensordot(be, a, b, 1)), - _ref_tropical_tensordot(anp, bnp, 1), - ) - np.testing.assert_allclose( - np.array(alg.einsum(be, "ab,bc->ac", a, b)), - _ref_tropical_einsum("ab,bc->ac", anp, bnp), - ) - - -def test_tropical_einsum_single_tensor_transpose_ok(): - # single tensor with NO repeated indices -> standard transpose is correct - be = _be() - rng = np.random.default_rng(5) - anp = rng.normal(size=(3, 4)) - a = be.cast(be.convert_to_tensor(anp), "float64") - got = np.array(_tropical_einsum(be, "ab->ba", a)) - np.testing.assert_allclose(got, np.array(be.einsum("ab->ba", a))) - - -def test_tropical_einsum_single_tensor_trace(): - # "ii->" : tropical trace = max of the diagonal - be = _be() - rng = np.random.default_rng(6) - square_a = rng.normal(size=(3, 3)) - a = be.cast(be.convert_to_tensor(square_a), "float64") - got = float(np.array(_tropical_einsum(be, "ii->", a))) - assert np.isclose(got, np.max(np.diag(square_a))) - - -def test_tropical_einsum_single_tensor_diagonal(): - # "ii->i" : tropical diagonal gather (repeated index kept, no reduction) - be = _be() - rng = np.random.default_rng(7) - square_a = rng.normal(size=(4, 4)) - a = be.cast(be.convert_to_tensor(square_a), "float64") - got = np.array(_tropical_einsum(be, "ii->i", a)) - np.testing.assert_allclose(got, np.diag(square_a)) - - -def test_tropical_einsum_single_tensor_reduce(): - # "ab->a" : single-tensor axis reduction (no repeat) must be tropical max, not sum - be = _be() - rng = np.random.default_rng(8) - anp = rng.normal(size=(3, 5)) - a = be.cast(be.convert_to_tensor(anp), "float64") - got = np.array(_tropical_einsum(be, "ab->a", a)) - np.testing.assert_allclose(got, np.max(anp, axis=1)) - - -def test_tropical_einsum_single_tensor_partial_trace(): - # "iij->j" : repeated index reduced (diagonal over i, then max) - be = _be() - rng = np.random.default_rng(9) - anp = rng.normal(size=(4, 4, 3)) # axes: i, i, j - a = be.cast(be.convert_to_tensor(anp), "float64") - got = np.array(_tropical_einsum(be, "iij->j", a)) - ref = np.max(np.array([anp[i, i, :] for i in range(4)]), axis=0) - np.testing.assert_allclose(got, ref) - - -def test_tropical_einsum_intra_operand_repeat_first(): - # "iij,jk->ik": first operand has a repeated index i (diagonal gather), then - # ordinary pairwise tropical contraction over j. Spec §4.5 step 1. - be = _be() - rng = np.random.default_rng(20) - anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 5)) - a = be.cast(be.convert_to_tensor(anp), "float64") - b = be.cast(be.convert_to_tensor(bnp), "float64") - got = np.array(_tropical_einsum(be, "iij,jk->ik", a, b)) - np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jk->ik", anp, bnp)) - - -def test_tropical_einsum_intra_operand_repeat_contracted(): - # "iij,jk->k": the repeated index i is neither in the output nor shared with - # b -> diagonal gather of i, then tropical max over both i and j. - be = _be() - rng = np.random.default_rng(21) - anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 5)) - a = be.cast(be.convert_to_tensor(anp), "float64") - b = be.cast(be.convert_to_tensor(bnp), "float64") - got = np.array(_tropical_einsum(be, "iij,jk->k", a, b)) - np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jk->k", anp, bnp)) - - -def test_tropical_einsum_both_operands_repeat(): - # "iij,jj->ij": both operands have intra-operand repeats; no contraction axis. - be = _be() - rng = np.random.default_rng(22) - anp, bnp = rng.normal(size=(3, 3, 4)), rng.normal(size=(4, 4)) - a = be.cast(be.convert_to_tensor(anp), "float64") - b = be.cast(be.convert_to_tensor(bnp), "float64") - got = np.array(_tropical_einsum(be, "iij,jj->ij", a, b)) - np.testing.assert_allclose(got, _ref_tropical_einsum("iij,jj->ij", anp, bnp)) - - -def test_tropical_context_changes_result(): - # max-plus "ab,b->a" = max_b (A[a,b] + B[b]) differs from standard sum-product - be = _be() - anp = np.array([[1.0, 5.0], [3.0, 2.0]]) - bnp = np.array([10.0, 0.0]) - a = be.cast(be.convert_to_tensor(anp), "float64") - b = be.cast(be.convert_to_tensor(bnp), "float64") - na, nb = tn.Node(a), tn.Node(b) - tn.connect(na[1], nb[0]) - expected = np.array([max(1 + 10, 5 + 0), max(3 + 10, 2 + 0)]) # [11, 13] - with tropical(): - got = np.array(cons.contractor([na, nb], output_edge_order=[na[0]]).tensor) - np.testing.assert_allclose(got, expected) - - -def test_preprocessing_does_not_corrupt_tropical(): - # Regression: default cons.contractor (greedy + preprocessing=True) calls - # _merge_single_gates for hyperedge-free >=5-node networks, merging single-gate - # nodes via STANDARD tn.contract_parallel BEFORE the algebraic path runs. For a - # tropical TN of regular tn.Nodes (no CopyNode -> no hyperedge -> preprocessing - # not skipped), that merge would corrupt the max-plus result with sum-product. - # The in-source _merge_single_gates guard skips the merge under a non-standard algebra. - import itertools - - be = _be() - rng = np.random.default_rng(0) - n0 = rng.normal(size=(2,)) - n1 = rng.normal(size=(2, 2)) - n2 = rng.normal(size=(2, 2)) - n3 = rng.normal(size=(2, 2)) - n4 = rng.normal(size=(2,)) - na = tn.Node(be.cast(be.convert_to_tensor(n0), "float64")) - nb = tn.Node(be.cast(be.convert_to_tensor(n1), "float64")) - nc = tn.Node(be.cast(be.convert_to_tensor(n2), "float64")) - nd = tn.Node(be.cast(be.convert_to_tensor(n3), "float64")) - ne = tn.Node(be.cast(be.convert_to_tensor(n4), "float64")) - tn.connect(na[0], nb[0]) - tn.connect(nb[1], nc[0]) - tn.connect(nc[1], nd[0]) - tn.connect(nd[1], ne[0]) - nodes = [na, nb, nc, nd, ne] - assert len(nodes) >= 5 - - ref = -np.inf - for s0, s1, s2, s3 in itertools.product(range(2), repeat=4): - val = n0[s0] + n1[s0, s1] + n2[s1, s2] + n3[s2, s3] + n4[s3] - ref = max(ref, val) - - with tropical(): - got = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() - - np.testing.assert_allclose(got, ref, atol=1e-6) - - -def test_tropical_public_api_surface(): - import applications.tropical_algebra as ex - - for name in [ - "MaxPlusAlgebra", - "MaxPlusTrackingAlgebra", - "CountingTropicalAlgebra", - "tropical", - "counting_tropical", - "recover_configuration", - "split_energy_count", - ]: - assert hasattr(ex, name), name diff --git a/tests/test_tropical_config.py b/tests/test_tropical_config.py deleted file mode 100644 index 13b64b79..00000000 --- a/tests/test_tropical_config.py +++ /dev/null @@ -1,246 +0,0 @@ -"""Task A1 (SPIKE) validation: argmax-tracking + tree-backtracking recovers an -optimal configuration of a tiny tropical (max, +) Ising contraction. - -The assertion is energy-based: the configuration returned by -``recover_configuration()`` must achieve the brute-force optimum energy. It need -not equal a *specific* brute-force argmin when the ground state is degenerate, -but its energy must be optimal. -""" - -import itertools - -import numpy as np -import pytest -import tensornetwork as tn - -import tensorcircuit as tc -import tensorcircuit.cons as cons -from applications.tropical_algebra import ( - tropical, - recover_configuration, - get_recorded_topology, - _td_backtrack_step, -) -from tests._tropical_test_utils import build_ising_tn - - -def _brute_cfg(n, edges, j_vals, h): - """Brute-force min energy over spin configs; returns (best_e, set_of_optimal_cfgs). - - cfg in {0,1}^n; spin s = 1 - 2*cfg (cfg=0 -> s=+1). E = -sum J s_i s_j - sum h s_i. - """ - best_e = None - best_cfgs = set() - for cfg in itertools.product([0, 1], repeat=n): - e = 0.0 - for (i, j), jij in zip(edges, j_vals): - e -= jij * (1 - 2 * cfg[i]) * (1 - 2 * cfg[j]) - for i, hi in zip(range(n), h): - e -= hi * (1 - 2 * cfg[i]) - if best_e is None or e < best_e - 1e-9: - best_e = e - best_cfgs = {cfg} - elif abs(e - best_e) < 1e-9: - best_cfgs.add(cfg) - return best_e, best_cfgs - - -def _energy_of(assignment, input_sets, raw_tensors): - """Max-plus total for a symbol->value assignment = -E for that config.""" - total = 0.0 - for term, t in zip(input_sets, raw_tensors): - idx = tuple(int(assignment[c]) for c in term) - total += float(np.asarray(t)[idx]) - return total - - -def _contract_tropical_track(nodes): - """Contract under tracking tropical; return (value, config).""" - with tropical(track=True): - val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() - cfg = recover_configuration() - return val, cfg - - -def test_config_recovery_tiny_ising_unique(): - # 2-spin Ising: spin0 --J-- spin1, fields h0, h1 -> unique ground state. - spins = [0, 1] - edges = [(0, 1)] - J = [0.7] - h = [0.3, -0.2] - best_e, best_cfgs = _brute_cfg(len(spins), edges, J, h) - - nodes = build_ising_tn(spins, edges, J, h) - val, cfg = _contract_tropical_track(nodes) - - # contraction returns max_cfg(-E) = -E_ground - np.testing.assert_allclose(val, -best_e, atol=1e-6) - - _tree, input_sets, raw_tensors = get_recorded_topology() - assert input_sets is not None and raw_tensors is not None - # every index label got a value - all_labels = set().union(*[set(t) for t in input_sets]) - assert set(cfg.keys()) == all_labels - - recovered_total = _energy_of(cfg, input_sets, raw_tensors) - # the recovered config's max-plus total must equal the contraction value - # and hence -best_e (optimal energy). - np.testing.assert_allclose(recovered_total, val, atol=1e-6) - np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) - - # unique ground state: the recovered (symbol) config must map to one of the - # brute-force optimal cfgs. Map symbols -> per-spin cfg via the field tensors: - # the field tensor for spin i has a single symbol whose value is cfg[i]. - # Recover that mapping from the topology (field tensors are the rank-1 terms). - spin_cfg = {} - for term, _t in zip(input_sets, raw_tensors): - if len(term) == 1: # field tensor -> one spin - spin_cfg[len(spin_cfg)] = int(cfg[term[0]]) - assert tuple(spin_cfg[i] for i in range(len(spins))) in best_cfgs - - -def test_config_recovery_tiny_ising_degenerate(): - # Degenerate ground state (Z2 spin-flip symmetry): ferromagnetic bond, no field. - # Any recovered config must still be optimal (energy-based assertion). - spins = [0, 1] - edges = [(0, 1)] - J = [1.0] - h = [0.0, 0.0] - best_e, best_cfgs = _brute_cfg(len(spins), edges, J, h) - assert len(best_cfgs) == 2 # all-up and all-down - - nodes = build_ising_tn(spins, edges, J, h) - val, cfg = _contract_tropical_track(nodes) - - np.testing.assert_allclose(val, -best_e, atol=1e-6) - _tree, input_sets, raw_tensors = get_recorded_topology() - recovered_total = _energy_of(cfg, input_sets, raw_tensors) - np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) - - spin_cfg = {} - for term, _t in zip(input_sets, raw_tensors): - if len(term) == 1: - spin_cfg[len(spin_cfg)] = int(cfg[term[0]]) - assert tuple(spin_cfg[i] for i in range(len(spins))) in best_cfgs - - -def test_config_recovery_three_ring(): - # Slightly bigger: 3-spin open chain -> exercises a 3rd pairwise step. - spins = [0, 1, 2] - edges = [(0, 1), (1, 2)] - J = [0.5, -0.8] - h = [0.1, 0.2, -0.3] - best_e, _best_cfgs = _brute_cfg(len(spins), edges, J, h) - - nodes = build_ising_tn(spins, edges, J, h) - val, cfg = _contract_tropical_track(nodes) - np.testing.assert_allclose(val, -best_e, atol=1e-6) - - _tree, input_sets, raw_tensors = get_recorded_topology() - recovered_total = _energy_of(cfg, input_sets, raw_tensors) - np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) - - -def test_config_recovery_non_scalar_raises(): - """``recover_configuration()`` is scalar-only: a contraction whose - root has dangling/free output indices (non-scalar result) must raise - ``NotImplementedError`` rather than return a (wrong) configuration. - - Builds a max-plus matrix-vector product - ``result[i] = max_k(A[i,k] + B[k])`` (one free index ``i``, one contracted - index ``k``) -- a non-scalar root -- and asserts recovery raises. The - non-scalar backtracking wiring has known ordering bugs (tree.output order - vs result-shape order; output-label values lost in the einsum branch), so it - is gated instead of shipping wrong answers. - """ - be = tc.backend - a_np = np.array( - [[1.0, 4.0, 2.0], [3.0, 1.0, 5.0], [0.0, 2.0, 1.0]] - ) # A[i,k], shape (3,3) - b_np = np.array([0.5, 1.0, -0.5]) # B[k], shape (3,) - - a_node = tn.Node(be.cast(be.convert_to_tensor(a_np), "float64")) - b_node = tn.Node(be.cast(be.convert_to_tensor(b_np), "float64")) - tn.connect(a_node[1], b_node[0]) # contract k; leave i (a_node[0]) dangling - nodes = [a_node, b_node] - - with tropical(track=True): - res = cons.contractor(nodes, output_edge_order=[a_node[0]]) - # sanity: the contraction itself is fine and non-scalar (shape (3,)) - assert np.array(res.tensor).shape == (3,) - with pytest.raises(NotImplementedError): - recover_configuration() - - -class _StubTree: - """Minimal stand-in for a cotengra ``ContractionTree`` exposing only the - three methods that ``_td_backtrack_step`` reads (``get_inds``, - ``get_tensordot_axes``, ``get_tensordot_perm``). Used to feed the backtrack - step a synthetic non-``None`` perm without spinning up a real contraction. - """ - - def __init__(self, inds, axes, perm): - self._inds = inds - self._axes = axes - self._perm = perm - - def get_inds(self, node): - return self._inds[node] - - def get_tensordot_axes(self, node): - return self._axes[node] - - def get_tensordot_perm(self, node): - return self._perm[node] - - -def test_td_backtrack_step_perm_inversion(): - """Synthetic-perm unit test for ``_td_backtrack_step`` (the I1 guard). - - The scalar 5-ring canary contracts with all-``None`` ``get_tensordot_perm`` - values, so the perm!=None inversion branch -- mapping a canonical - (``get_inds(p)``) position to the algebra-natural (``l_free + r_free``) - order via ``td_pos[perm[i]] = p_pos[i]`` -- is NEVER exercised end-to-end. - This feeds the backtrack step a synthetic 3-cycle perm (non-self-inverse, so - a wrong inversion *direction* is caught, unlike a transposition) and asserts - the recovered per-axis positions are exactly the inverse of the perm. - - Guards the inversion logic for future non-scalar support and against cotengra - convention changes. - """ - # Geometry: l = 'ab' (a free, b contracted), r = 'bcd' (b contracted, c,d - # free). l_axes=[1] (b is axis 1 of l), r_axes=[0] (b is axis 0 of r). - # Algebra-natural output order = l_free + r_free = 'acd'. - # Canonical (get_inds(p)) order chosen as 'dac' (a 3-cycle of 'acd'), so the - # perm is non-trivial and non-self-inverse: - # perm[i] = 'acd'.find('dac'[i]) = (2, 0, 1). - l_node, r_node, p_node = "l", "r", "p" - tree = _StubTree( - inds={l_node: "ab", r_node: "bcd", p_node: "dac"}, - axes={p_node: ([1], [0])}, - perm={p_node: (2, 0, 1)}, - ) - # argmax lives in algebra-natural 'acd' order, shape (dim_a, dim_c, dim_d). - # Place flattened contracted index 3 at td_pos (a=1, c=2, d=0); with - # contract_dims=(4,) -> contracted label 'b' value 3. - argmax = np.zeros((2, 3, 2), dtype=np.int64) - argmax[1, 2, 0] = 3 - rec = {"kind": "td", "argmax": argmax, "contract_dims": (4,)} - # Canonical p_pos in 'dac' order: d=0, a=1, c=2. - p_pos = (0, 1, 2) - assignment = {} - l_pos, r_pos = _td_backtrack_step( - tree, p_node, l_node, r_node, rec, p_pos, assignment - ) - - # Inversion applied: canonical 'dac'(0,1,2) -> algebra-natural 'acd'(1,2,0), - # i.e. td_pos = (a=1, c=2, d=0); argmax[1,2,0]=3 -> 'b'=3. - assert assignment["b"] == 3 - # l_pos in 'ab' order: a=1 (free via inversion), b=3 (contracted). - assert l_pos == (1, 3) - # r_pos in 'bcd' order: b=3 (contracted), c=2 (free), d=0 (free). - assert r_pos == (3, 2, 0) - - -if __name__ == "__main__": - pytest.main([__file__, "-q"]) diff --git a/tests/test_tropical_counting.py b/tests/test_tropical_counting.py deleted file mode 100644 index 84bbce19..00000000 --- a/tests/test_tropical_counting.py +++ /dev/null @@ -1,364 +0,0 @@ -import numpy as np -import tensorcircuit as tc -from applications.tropical_algebra import ( - _counting_tensordot, - split_energy_count, -) - - -def _be(): - return tc.backend - - -def _ref_counting_tensordot(anp, bnp, axes, eps=1e-9): - # brute-force (max, degeneracy) tensordot on energy-only inputs (count starts at 1) - if isinstance(axes, int): - a_axes = list(range(anp.ndim - axes, anp.ndim)) - b_axes = list(range(0, axes)) - else: - a_axes, b_axes = list(axes[0]), list(axes[1]) - a_free = [i for i in range(anp.ndim) if i not in a_axes] - b_free = [i for i in range(bnp.ndim) if i not in b_axes] - at = np.transpose(anp, a_free + a_axes) - bt = np.transpose(bnp, b_axes + b_free) - a_fs = [anp.shape[i] for i in a_free] - b_fs = [bnp.shape[i] for i in b_free] - k = int(np.prod([anp.shape[i] for i in a_axes]) or 1) - A = int(np.prod(a_fs) or 1) - b_n = int(np.prod(b_fs) or 1) - a2 = at.reshape(A, k) - b2 = bt.reshape(k, b_n) - out_e = np.full((A, b_n), -np.inf) - out_n = np.zeros((A, b_n)) - for i in range(A): - for j in range(b_n): - s = a2[i, :] + b2[:, j] - mx = np.max(s) - out_e[i, j] = mx - out_n[i, j] = int(np.sum(np.abs(s - mx) < eps)) # degeneracy = #max-tied - return out_e.reshape(a_fs + b_fs), out_n.reshape(a_fs + b_fs) - - -def test_counting_tensordot_matrix(): - be = _be() - rng = np.random.default_rng(10) - anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) - # stacked inputs: [x, n], counts init to 1 - a = be.cast( - be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], axis=-1)), "float64" - ) - b = be.cast( - be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], axis=-1)), "float64" - ) - got = np.array(_counting_tensordot(be, a, b, axes=1)) - ref_e, ref_n = _ref_counting_tensordot(anp, bnp, 1) - np.testing.assert_allclose(got[..., 0], ref_e) - np.testing.assert_allclose(got[..., 1], ref_n) - - -def test_split_energy_count(): - be = _be() - arr = np.array([[1.0, 2.0], [3.0, 4.0]]) - stacked = be.cast( - be.convert_to_tensor(np.stack([arr, arr * 2], axis=-1)), "float64" - ) - e, n = split_energy_count(stacked) - np.testing.assert_allclose(np.array(e), arr) - np.testing.assert_allclose(np.array(n), arr * 2) - - -# --- Task B2: counting einsum (hyperedge / copy-node aware) --- - -from applications.tropical_algebra import ( # noqa: E402 - _counting_einsum, -) - - -def _ref_counting_einsum(eq, a, b, eps=1e-9): - """Brute-force (energy, degeneracy) 2-operand einsum on energy-only inputs.""" - import itertools - - lhs, rhs = eq.split("->") - ia, ib = lhs.split(",") - sizes = {} - for s, t in zip([ia, ib], [a, b]): - for c, dim in zip(s, t.shape): - sizes[c] = dim - out_e = np.full([sizes[c] for c in rhs], -np.inf) - out_n = np.zeros([sizes[c] for c in rhs]) - allc = list(dict.fromkeys(list(ia) + list(ib))) - for combo in itertools.product(*[range(sizes[c]) for c in allc]): - env = dict(zip(allc, combo)) - e = a[tuple(env[c] for c in ia)] + b[tuple(env[c] for c in ib)] - oidx = tuple(env[c] for c in rhs) - if e > out_e[oidx] + eps: - out_e[oidx] = e - out_n[oidx] = 1 - elif abs(e - out_e[oidx]) < eps: - out_n[oidx] += 1 - return out_e, out_n - - -def test_counting_einsum_pair(): - be = _be() - rng = np.random.default_rng(11) - anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) - a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") - b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") - got = np.array(_counting_einsum(be, "ab,bc->ac", a, b)) - ref_e, ref_n = _ref_counting_einsum("ab,bc->ac", anp, bnp) - np.testing.assert_allclose(got[..., 0], ref_e) - np.testing.assert_allclose(got[..., 1], ref_n) - - -def test_counting_einsum_hyperedge(): - # shared index 'a' (batch) — copy-node pair shape - be = _be() - rng = np.random.default_rng(12) - anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) - a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") - b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") - got = np.array(_counting_einsum(be, "ab,ac->abc", a, b)) - ref_e, ref_n = _ref_counting_einsum("ab,ac->abc", anp, bnp) - np.testing.assert_allclose(got[..., 0], ref_e) - np.testing.assert_allclose(got[..., 1], ref_n) - - -def test_counting_einsum_tie_degeneracy(): - # Symmetric inputs over the contracted axis -> a clean tie (degeneracy > 1). - be = _be() - # a: shape (1, 2) with equal values on the contracted axis; b: shape (2, 1). - anp = np.array([[1.0, 1.0]]) # both k=0 and k=1 contribute the same energy - bnp = np.array([[2.0], [2.0]]) - a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") - b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") - got = np.array(_counting_einsum(be, "ab,bc->ac", a, b)) - ref_e, ref_n = _ref_counting_einsum("ab,bc->ac", anp, bnp) - np.testing.assert_allclose(got[..., 0], ref_e) - np.testing.assert_allclose(got[..., 1], ref_n) - # explicit assertion: energy 3.0 achieved by both k values -> degeneracy 2 - assert ref_n[0, 0] == 2 - np.testing.assert_allclose(got[0, 0, 0], 3.0) - np.testing.assert_allclose(got[0, 0, 1], 2.0) - - -def test_counting_einsum_single_operand_diagonal(): - # eq "aa->a": diagonal, then no contraction -> energy=diag(e), count=diag(n) - import tensorcircuit.backends.numpy_backend as nb - - be = nb.NumpyBackend() - e = np.array([[1.0, 5.0], [3.0, 2.0]]) - n = np.array([[1.0, 2.0], [1.0, 1.0]]) - a = np.stack([e, n], axis=-1) - out = _counting_einsum(be, "aa->a", a) - np.testing.assert_allclose(out[..., 0], np.diag(e)) - np.testing.assert_allclose(out[..., 1], np.diag(n)) - - -def test_counting_einsum_single_operand_reduce(): - # eq "ab->a": max over b of energy, tie-sum count - import tensorcircuit.backends.numpy_backend as nb - - be = nb.NumpyBackend() - e = np.array([[1.0, 5.0], [3.0, 3.0]]) - n = np.array([[1.0, 2.0], [1.0, 4.0]]) - a = np.stack([e, n], axis=-1) - out = _counting_einsum(be, "ab->a", a) - np.testing.assert_allclose(out[..., 0], [5.0, 3.0]) - np.testing.assert_allclose(out[..., 1], [2.0, 5.0]) # ties: b=0,1 both 3 -> 1+4 - - -def test_counting_einsum_two_operand_intra_repeat(): - # eq "aa,b->ab": operand a has repeated index - import tensorcircuit.backends.numpy_backend as nb - - be = nb.NumpyBackend() - e_a = np.array([[2.0, 1.0], [1.0, 2.0]]) - n_a = np.ones((2, 2)) - a = np.stack([e_a, n_a], axis=-1) - e_b = np.array([0.0, 0.0]) - n_b = np.array([1.0, 1.0]) - b = np.stack([e_b, n_b], axis=-1) - out = _counting_einsum(be, "aa,b->ab", a, b) - # diagonal of a is [2,2]; outer with b=[0,0] -> [[2,2],[2,2]], counts follow - np.testing.assert_allclose(out[..., 0], [[2, 2], [2, 2]]) - # count stream: diag(n_a)=[1,1] outer n_b=[1,1] -> [[1,1],[1,1]] - np.testing.assert_allclose(out[..., 1], [[1, 1], [1, 1]]) - - -def test_counting_einsum_multi_axis(): - # Multi-axis contraction "ab,ac->c": shared index a (contracted) plus a free - # index b of operand a (also summed out) -> two axes reduced. Compared to a - # brute-force (energy, degeneracy) reference. - be = _be() - rng = np.random.default_rng(14) - anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(3, 5)) - a = be.cast(be.convert_to_tensor(np.stack([anp, np.ones_like(anp)], -1)), "float64") - b = be.cast(be.convert_to_tensor(np.stack([bnp, np.ones_like(bnp)], -1)), "float64") - got = np.array(_counting_einsum(be, "ab,ac->c", a, b)) - ref_e, ref_n = _ref_counting_einsum("ab,ac->c", anp, bnp) - np.testing.assert_allclose(got[..., 0], ref_e) - np.testing.assert_allclose(got[..., 1], ref_n) - - -# --- Task 8: degeneracy side-channel + standard-contraction clear discipline --- - - -def test_degeneracy_none_after_standard_contraction(): - """A counting contraction stashes count=deg into the aux side-channel; a - subsequent STANDARD contraction must clear it (unconditional clear in - ``cons._algebraic_base_contraction``), so ``degeneracy()`` returns None - rather than a stale count from the earlier counting contraction. - """ - import opt_einsum - import tensornetwork as tn - import tensorcircuit.cons as cons - import applications.tropical_algebra as tr - from tests._tropical_test_utils import ( - build_ising_tn, - brute_force_energy_and_degeneracy, - ) - - spins = [0, 1, 2, 3, 4] - edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] # 5-ring - J = [1.0, 1.0, 1.0, 1.0, 1.0] # ferromagnetic - h = [0.0, 0.0, 0.0, 0.0, 0.0] # zero field -> Z2 spin-flip symmetry, g=2 - expected_e, expected_n = brute_force_energy_and_degeneracy(len(spins), edges, J, h) - assert ( - expected_n == 2 - ) # guard: this instance MUST be degenerate or the test is moot - nodes = build_ising_tn(spins, edges, J, h) - - with tr.counting_tropical(): - node = cons._algebraic_base_contraction( - nodes, - algorithm=opt_einsum.paths.dynamic_programming, - output_edge_order=[], - ignore_edge_order=True, - ) - # counting decode stashed the count; degeneracy() reads it (valid here) - np.testing.assert_allclose(np.array(tr.degeneracy()), expected_n) - np.testing.assert_allclose(np.array(node.tensor), -expected_e, atol=1e-6) - - # exit -> StandardAlgebra restored. A standard contraction must clear aux - # (the unconditional _stash_aux_outputs({}) at the top of every contraction), - # so a stale count cannot leak out of the counting block. - rng = np.random.default_rng(0) - a = tn.Node(rng.standard_normal((2, 3)).astype(np.complex64)) - b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) - tn.connect(a[1], b[0]) - cons._algebraic_base_contraction( - [a, b], - algorithm=opt_einsum.paths.dynamic_programming, - output_edge_order=[a[0], b[1]], - ) - assert tr.degeneracy() is None # standard contraction cleared aux - - -# --- Task 10: non-scalar counting canary (aux co-indexed with output_edge_order) --- - - -def test_nonscalar_counting_energy_and_degeneracy_per_output(): - """A counting contraction with ONE free (dangling) spin returns energy AND - degeneracy PER output configuration, with the count aux co-indexed with the - energy via ``output_edge_order``. - - Validates two things at once: - (1) non-scalar counting works at all -- ``decode`` strips the trailing count - axis so ``tn.Node`` wraps a rank-correct energy tensor (rank == #free); - (2) the aux-reorder logic (Task 5's ``output_edge_order`` permutation applied - to aux) aligns count with energy. If the perm were wrong, count would be - permuted differently from energy and the per-output assertion would fail. - - Brute force: for each value v in {0,1} of the free spin (cfg=0 -> s=+1, - cfg=1 -> s=-1), enumerate the other spins, find min E and its degeneracy. - """ - import opt_einsum - import applications.tropical_algebra as tr - import tensorcircuit.cons as cons - from tests._tropical_test_utils import ( - build_ring_with_free_spin, - brute_nonscalar_counting, - ) - - nodes, free_edge, J, h = build_ring_with_free_spin() - expected_e, expected_n = brute_nonscalar_counting(nodes, free_edge, J, h) - - with tr.counting_tropical(): - node = cons._algebraic_base_contraction( - nodes, - algorithm=opt_einsum.paths.dynamic_programming, - output_edge_order=[free_edge], - ) - E = np.asarray(node.tensor) - N = np.asarray(tr.degeneracy()) - - np.testing.assert_allclose(E, expected_e) - np.testing.assert_allclose(N, expected_n) - # guard: the two output configs must differ in BOTH energy and count, - # otherwise a permuted aux could still pass the assertion by accident. - assert E[0] != E[1] or N[0] != N[1] - - -# --- Task 10 follow-up: 2-free-spin perm-direction canary --- - - -def test_nonscalar_counting_two_free_spins_reversed_order(): - """A counting contraction with TWO free (dangling) spins where - ``output_edge_order`` is INTENTIONALLY REVERSED relative to the dangling-edge - order, so the aux-reorder permutation in ``cons._algebraic_base_contraction`` - (``perm = [dangling_edges.index(e) for e in order]``) is non-trivial - (``(1, 0)`` swap, not the identity ``(0,)`` exercised by the one-free-spin - canary). - - Validates that the SAME permutation applied to the energy tensor (via - ``final_node.reorder_edges``) is also applied to the aux degeneracy tensor - (via ``kbe.transpose(v, tuple(perm))``). If the aux perm were ELIDED, the - energy ``E`` would come out in the requested ``[b, a]`` order but the count - ``N`` would be left in the dangling ``[a, b]`` order, so - ``assert_allclose(N, expected_n_ab.T)`` would FAIL (the brute-force result is - asymmetric under transpose by parameter choice -- see the guard below). - - The energy ``E`` itself is also checked against the transposed brute force, - pinning down which order the contraction actually produced. - """ - import opt_einsum - import applications.tropical_algebra as tr - import tensorcircuit.cons as cons - from tests._tropical_test_utils import ( - build_ring_with_two_free_spins, - brute_nonscalar_counting_2d, - ) - - nodes, edge_a, edge_b, J, h = build_ring_with_two_free_spins() - # brute force in the natural (a, b) order: - expected_e_ab, expected_n_ab = brute_nonscalar_counting_2d( - nodes, edge_a, edge_b, J, h - ) - # guard: the brute-force result MUST be asymmetric under transpose, otherwise - # a wrong-order aux could still pass ``assert_allclose(N, expected_n_ab.T)`` - # by accident. The default parameters give E=[[5,7],[5,-1]], N=[[1,2],[1,1]], - # both asymmetric. - assert not np.allclose(expected_e_ab, expected_e_ab.T), ( - "weak canary: brute-force energy is symmetric under transpose, so the " - "perm direction is not load-bearing for E" - ) - assert not np.allclose(expected_n_ab, expected_n_ab.T), ( - "weak canary: brute-force degeneracy is symmetric under transpose, so " - "the perm direction is not load-bearing for N" - ) - - # contract with output_edge_order REVERSED -> [edge_b, edge_a]; under - # sorted_edges the dangling order is [edge_a, edge_b], so the perm is (1, 0). - with tr.counting_tropical(): - node = cons._algebraic_base_contraction( - nodes, - algorithm=opt_einsum.paths.dynamic_programming, - output_edge_order=[edge_b, edge_a], - ) - E = np.asarray(node.tensor) # in [b, a] order (reorder_edges applied) - N = np.asarray(tr.degeneracy()) # MUST also be in [b, a] order if the - # # aux perm is applied correctly. - # expected in [b, a] order = transpose of the (a, b) brute force: - np.testing.assert_allclose(E, expected_e_ab.T) - np.testing.assert_allclose(N, expected_n_ab.T) diff --git a/tests/test_tropical_example.py b/tests/test_tropical_example.py deleted file mode 100644 index 40507345..00000000 --- a/tests/test_tropical_example.py +++ /dev/null @@ -1,39 +0,0 @@ -import numpy as np -import tensornetwork as tn -import tensorcircuit as tc -import tensorcircuit.cons as cons -from applications.tropical_algebra import ( - MaxPlusAlgebra, - tropical, - recover_configuration, -) - - -def test_example_maxplus_tensordot_matches_brute(): - be = tc.backend - rng = np.random.default_rng(0) - anp, bnp = rng.normal(size=(3, 4)), rng.normal(size=(4, 5)) - a = be.cast(be.convert_to_tensor(anp), "float64") - b = be.cast(be.convert_to_tensor(bnp), "float64") - got = np.array(MaxPlusAlgebra().tensordot(be, a, b, 1)) - # brute max-plus: Y[i,j] = max_k a[i,k] + b[k,j] - ref = np.max(anp[:, :, None] + bnp[None, :, :], axis=1) - np.testing.assert_allclose(got, ref) - - -def test_example_recover_configuration_on_tiny_ising(): - # 3-node chain (vector-matrix-vector) -> recover_configuration round-trips - be = tc.backend - n = tn.Node(be.cast(be.convert_to_tensor(np.array([0.0, 0.0])), "float64")) - e = tn.Node( - be.cast(be.convert_to_tensor(np.array([[1.0, -1.0], [-1.0, 1.0]])), "float64") - ) - m = tn.Node(be.cast(be.convert_to_tensor(np.array([0.5, -0.5])), "float64")) - tn.connect(n[0], e[0]) - tn.connect(e[1], m[0]) - nodes = [n, e, m] - with tropical(track=True): - float(cons.contractor(nodes, output_edge_order=[]).tensor) - cfg = recover_configuration() - assert isinstance(cfg, dict) - assert len(cfg) >= 1 diff --git a/tests/test_tropical_ising.py b/tests/test_tropical_ising.py deleted file mode 100644 index 14ced604..00000000 --- a/tests/test_tropical_ising.py +++ /dev/null @@ -1,223 +0,0 @@ -import itertools -import numpy as np -import tensorcircuit.cons as cons -from applications.tropical_algebra import tropical, counting_tropical -from tests._tropical_test_utils import ( - build_ising_tn, - brute_force_energy, - brute_force_energy_and_degeneracy, -) - - -def test_ising_ring_ground_state_matches_bruteforce(): - spins = [0, 1, 2, 3, 4] - edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] # 5-ring - J = [1.0, -1.0, 1.0, 1.0, -1.0] - h = [0.5, -0.3, 0.2, 0.0, 0.4] - nodes = build_ising_tn(spins, edges, J, h) - - e_ground = brute_force_energy(spins, edges, J, h) - - with tropical(): - val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() - - # contraction returns max_cfg(-E) = -E_ground - np.testing.assert_allclose(val, -e_ground, atol=1e-6) - - -# --- Task A2: configuration recovery (argmax backtracking) canary --- - - -def _ising_energy_of_cfg(cfg, edges, j_vals, h): - """Ising energy of a {0,1}-config: E = -sum J s_i s_j - sum h s_i (s=1-2cfg).""" - e = 0.0 - for (i, j), jij in zip(edges, j_vals): - si, sj = 1 - 2 * cfg[i], 1 - 2 * cfg[j] - e -= jij * si * sj - for i, hi in zip(range(len(h)), h): - e -= hi * (1 - 2 * cfg[i]) - return e - - -def _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors): - """Max-plus total of a symbol->value assignment = sum of tensor entries = - -E for that configuration. Mapping-independent optimality proof.""" - total = 0.0 - for term, t in zip(input_sets, raw_tensors): - idx = tuple(int(cfg_sym[c]) for c in term) - total += float(np.asarray(t)[idx]) - return total - - -def test_ising_config_recovery_five_ring(): - """Task A2 canary: recover an optimal spin configuration of the 5-spin Ising - ring under ``tropical(track=True)`` and verify its energy equals the - brute-force ground energy. - - The 5-ring mixes ``tensordot`` (ordinary pair) and ``einsum`` (hyperedge -- - a CopyNode hub of degree >= 2 puts its symbol on >= 3 regular nodes) - contraction steps, so this exercises both backtracking branches. Optimality - is asserted two ways: (1) the max-plus total of the recovered symbol - assignment equals the contraction value (== -E_ground) -- mapping- - independent; (2) translating the assignment to per-spin values via the - rank-1 field tensors and computing the Ising energy gives ``best_e``. - """ - from applications.tropical_algebra import ( - get_recorded_topology, - recover_configuration, - ) - - spins = [0, 1, 2, 3, 4] - edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] # 5-ring - J = [1.0, -1.0, 1.0, 1.0, -1.0] - h = [0.5, -0.3, 0.2, 0.0, 0.4] - - # brute-force ground energy over the 2^5 configs - best_e = None - for cfg in itertools.product([0, 1], repeat=len(spins)): - e = _ising_energy_of_cfg(cfg, edges, J, h) - if best_e is None or e < best_e - 1e-9: - best_e = e - - nodes = build_ising_tn(spins, edges, J, h) - with tropical(track=True): - val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() - cfg_sym = recover_configuration() - - # contraction value == -E_ground - np.testing.assert_allclose(val, -best_e, atol=1e-6) - - _tree, input_sets, raw_tensors = get_recorded_topology() - assert input_sets is not None and raw_tensors is not None - # every index label received a value - all_labels = set().union(*[set(t) for t in input_sets]) - assert set(cfg_sym.keys()) == all_labels - - # (1) mapping-independent optimality: max-plus total == contraction value - recovered_total = _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors) - np.testing.assert_allclose(recovered_total, val, atol=1e-6) - np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) - - # (2) translate symbol config -> per-spin cfg via rank-1 field tensors - # (Tv_i has a single index label; its value is cfg[i]). The field tensors - # appear in input_sets in spin order (regular_nodes sorted by _stable_id_, - # which follows creation order in build_ising_tn). - spin_cfg = {} - for term, _t in zip(input_sets, raw_tensors): - if len(term) == 1: - spin_cfg[len(spin_cfg)] = int(cfg_sym[term[0]]) - assert len(spin_cfg) == len(spins) - recovered_e = _ising_energy_of_cfg( - [spin_cfg[i] for i in range(len(spins))], edges, J, h - ) - np.testing.assert_allclose(recovered_e, best_e, atol=1e-6) - - -def test_ising_config_recovery_degenerate_ring(): - """Config recovery under degeneracy (Z2-symmetric ferromagnetic ring, zero - field): the recovered config must still be energy-optimal (it is one of the - degenerate ground states, selected by first-argument-wins tie-breaking).""" - from applications.tropical_algebra import ( - get_recorded_topology, - recover_configuration, - ) - - spins = [0, 1, 2, 3, 4] - edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] - J = [1.0, 1.0, 1.0, 1.0, 1.0] # ferromagnetic - h = [0.0, 0.0, 0.0, 0.0, 0.0] # zero field -> Z2 spin-flip symmetry, g=2 - - best_e = None - for cfg in itertools.product([0, 1], repeat=len(spins)): - e = _ising_energy_of_cfg(cfg, edges, J, h) - if best_e is None or e < best_e - 1e-9: - best_e = e - assert best_e == -5.0 # ferromagnetic ground state - - nodes = build_ising_tn(spins, edges, J, h) - with tropical(track=True): - val = np.array(cons.contractor(nodes, output_edge_order=[]).tensor).item() - cfg_sym = recover_configuration() - - np.testing.assert_allclose(val, -best_e, atol=1e-6) - _tree, input_sets, raw_tensors = get_recorded_topology() - recovered_total = _maxplus_total_of_assignment(cfg_sym, input_sets, raw_tensors) - np.testing.assert_allclose(recovered_total, -best_e, atol=1e-6) - - # the recovered spin config must be one of the two aligned ground states - spin_cfg = {} - for term, _t in zip(input_sets, raw_tensors): - if len(term) == 1: - spin_cfg[len(spin_cfg)] = int(cfg_sym[term[0]]) - recovered_e = _ising_energy_of_cfg( - [spin_cfg[i] for i in range(len(spins))], edges, J, h - ) - np.testing.assert_allclose(recovered_e, best_e, atol=1e-6) - - -# --- Task B3: counting tropical (energy*, degeneracy) canary --- - - -def test_ising_counting_ground_state_and_degeneracy(): - """Energy via ``node.tensor``; degeneracy via ``tr.degeneracy()`` side channel. - - Under the encode/decode design, ``CountingRepresentation.encode`` attaches - count=1 to each leaf, so the test builder emits PLAIN energy tensors (the - same ``build_ising_tn`` used by the max-plus tests). After contraction under - ``counting_tropical()``, the primary tensor is the energy (-E_ground) and - the count is stashed in the aux side-channel for ``degeneracy()`` to read. - """ - import applications.tropical_algebra as tr - - spins = [0, 1, 2, 3, 4] - edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] # 5-ring - J = [1.0, -1.0, 1.0, 1.0, -1.0] - h = [0.5, -0.3, 0.2, 0.0, 0.4] - # brute-force ground energy + degeneracy (s = 1 - 2*cfg; cfg=0 -> s=+1, cfg=1 -> s=-1) - best_e, deg = brute_force_energy_and_degeneracy(len(spins), edges, J, h) - - nodes = build_ising_tn(spins, edges, J, h) - - with counting_tropical(): - res = cons.contractor(nodes, output_edge_order=[], ignore_edge_order=True) - count = tr.degeneracy() - - energy = np.array(res.tensor) - # contraction returns max_cfg(-E) = -E_ground and the #cfgs achieving it. - np.testing.assert_allclose(energy, -best_e, atol=1e-6) - np.testing.assert_allclose(np.array(count), deg, atol=1e-6) - - -def test_ising_counting_degenerate_ring(): - """Degenerate ground state (g=2) canary for the counting tropical stream. - - The original ``test_ising_counting_ground_state_and_degeneracy`` instance is - tuned to a *unique* ground state (g=1), so a regression that always returns - ``count=1`` would pass undetected. A ferromagnetic ring with zero field has a - Z2 (global spin-flip) symmetry -> exactly 2 ground states (all-up and - all-down), so this case guards the degeneracy stream against such a regression. - """ - import applications.tropical_algebra as tr - - spins = [0, 1, 2, 3, 4] - edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 0)] # 5-ring - J = [1.0, 1.0, 1.0, 1.0, 1.0] # ferromagnetic - h = [0.0, 0.0, 0.0, 0.0, 0.0] # zero field -> Z2 spin-flip symmetry - # brute-force ground energy + degeneracy (s = 1 - 2*cfg; cfg=0 -> s=+1, cfg=1 -> s=-1) - best_e, deg = brute_force_energy_and_degeneracy(len(spins), edges, J, h) - - # sanity-check the derivation: 5 ferromagnetic bonds, all J=+1, h=0 -> - # ground state all-spins-aligned, E_ground = -sum(J) = -5 over exactly 2 - # configs (Z2-related). Contraction returns max_cfg(-E) = -E_ground = 5. - assert best_e == -5.0 - assert deg == 2 - - nodes = build_ising_tn(spins, edges, J, h) - - with counting_tropical(): - res = cons.contractor(nodes, output_edge_order=[], ignore_edge_order=True) - count = tr.degeneracy() - - energy = np.array(res.tensor) - np.testing.assert_allclose(energy, -best_e, atol=1e-6) # == 5 - np.testing.assert_allclose(np.array(count), deg, atol=1e-6) # == 2 From b218efd9a91e3f8bd531d94a1400e642303e5b79 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 00:44:22 +0800 Subject: [PATCH 05/16] refactor: remove algebra parameter from set_contractor --- applications/bcomplex32_algebra.py | 2 +- applications/tropical_algebra.py | 2 +- tensorcircuit/cons.py | 20 ++------ tensorcircuit/contraction_algebra/__init__.py | 4 +- tests/test_contraction_algebra.py | 51 +++++++++---------- 5 files changed, 32 insertions(+), 47 deletions(-) diff --git a/applications/bcomplex32_algebra.py b/applications/bcomplex32_algebra.py index 82798515..a04281ed 100644 --- a/applications/bcomplex32_algebra.py +++ b/applications/bcomplex32_algebra.py @@ -1,7 +1,7 @@ """complex pair-algebra — a reference APPLICATION of ContractionAlgebra. Pair repr: complex tensor = stack([re, im], axis=-1) of bf16. Contraction = 4 real -bf16 matmuls (4M). Activated via ``set_contractor(algebra=ComplexPairAlgebra())`` or +bf16 matmuls (4M). Activated via ``cons.set_contraction_algebra(ComplexPairAlgebra())`` or the ``bcomplex32()`` CM. encode/decode at the ``_algebraic_base_contraction`` boundary keep the pair axis off tn.Node (dodges the axis==edge wall). """ diff --git a/applications/tropical_algebra.py b/applications/tropical_algebra.py index a377351e..a2918385 100644 --- a/applications/tropical_algebra.py +++ b/applications/tropical_algebra.py @@ -4,7 +4,7 @@ The feature is ``tensorcircuit.contraction_algebra`` -- the ``ContractionAlgebra`` and ``Representation`` ABCs -- wired in-source into ``cons._base``, which routes any non-standard algebra through ``cons._algebraic_base_contraction`` (no monkey-patching). -A custom algebra is activated either via ``tc.set_contractor(algebra=...)`` or by the +A custom algebra is activated via ``tc.set_contraction_algebra(alg)`` or by the ``tropical()`` / ``counting_tropical()`` context managers below. This file is one complete algebra built on top of the ABCs -- an importable reference module, not a runnable demo (for that, see ``examples/tropical_ising.py``). Kept under diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 7d591659..4f4d6af5 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -1274,13 +1274,17 @@ def set_contractor( contraction_info: bool = False, debug_level: int = 0, use_primitives: Optional[bool] = None, - algebra: Optional[_ContractionAlgebra] = None, **kws: Any, ) -> Callable[..., Any]: """ To set runtime contractor of the tensornetwork for a better contraction path. For more information on the usage of contractor, please refer to independent tutorial. + 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 + are the recommended way to switch algebras for a block of code. + :param method: "auto", "greedy", "branch", "plain", "tng", "custom", "custom_stateful". Also supports shortcuts like "cotengra", "cotengra-30-64", "omeco", and "omeco-16-32". defaults to None ("auto") @@ -1295,20 +1299,6 @@ def set_contractor( :return: The new tensornetwork with its contractor set. :rtype: tn.Node """ - if algebra is not None: - if not isinstance(algebra, _StandardAlgebra): - use_primitives = True - if kws.get("preprocessing", False): - raise ValueError( - "preprocessing is incompatible with a non-standard " - "ContractionAlgebra (it contracts via native sum-product)" - ) - if kws.get("strip_exponent", False): - raise ValueError( - "strip_exponent is incompatible with a non-standard " - "ContractionAlgebra (its log-scaling assumes sum-product)" - ) - set_contraction_algebra(algebra) if not method: method = "greedy" # auto for small size fallbacks to dp, which has bug for now diff --git a/tensorcircuit/contraction_algebra/__init__.py b/tensorcircuit/contraction_algebra/__init__.py index d831e06a..b4c76cda 100644 --- a/tensorcircuit/contraction_algebra/__init__.py +++ b/tensorcircuit/contraction_algebra/__init__.py @@ -2,8 +2,8 @@ boundary representation, consulted by ``cons._algebraic_base_contraction``. Implement ``ContractionAlgebra`` (with a ``Representation``) and activate via -``cons.set_contraction_algebra(...)`` (or ``cons.set_contractor(algebra=...)``, -``cons.runtime_contractor(..., algebra=...)``). The in-source ``cons._base`` +``cons.set_contraction_algebra(...)`` -- the single entry-point for switching +contraction algebras. The in-source ``cons._base`` routes any non-standard algebra to ``_algebraic_base_contraction``, which runs encode -> algebra kernels -> decode; no monkey-patching is required. diff --git a/tests/test_contraction_algebra.py b/tests/test_contraction_algebra.py index c325a91c..068a8d67 100644 --- a/tests/test_contraction_algebra.py +++ b/tests/test_contraction_algebra.py @@ -41,7 +41,7 @@ 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 - # in-source via cons.set_contraction_algebra / set_contractor(algebra=...). + # in-source via cons.set_contraction_algebra. for name in [ "ContractionAlgebra", "StandardAlgebra", @@ -222,7 +222,7 @@ def einsum(self, be, eq, *ops): assert isinstance(BareAlgebra().representation, IdentityRepresentation) -# --- Task 2: cons.py algebra state + set_contractor(algebra=) --- +# --- Task 2: cons.py algebra state (set_contraction_algebra) --- class _NS(ContractionAlgebra): # non-standard stub for constraint tests @@ -236,48 +236,43 @@ def einsum(self, be, eq, *ops): return be.einsum(eq, *ops) -def test_set_contractor_accepts_algebra_kwarg(): +def test_set_contraction_algebra(): prev = cons.get_contraction_algebra() try: - cons.set_contractor("greedy", algebra=_NS()) + cons.set_contraction_algebra(_NS()) assert isinstance(cons.get_contraction_algebra(), _NS) finally: cons.set_contraction_algebra(prev) -def test_set_contractor_rejects_preprocessing_with_nonstandard_algebra(): - import pytest - - ns_algebra = _NS() - with pytest.raises(ValueError): - cons.set_contractor("greedy", algebra=ns_algebra, preprocessing=True) - - -def test_set_contractor_rejects_strip_exponent_with_nonstandard_algebra(): - import pytest - - ns_algebra = _NS() - with pytest.raises(ValueError): - cons.set_contractor("greedy", algebra=ns_algebra, strip_exponent=True) - - def test_runtime_contractor_restores_algebra(): prev = cons.get_contraction_algebra() - with cons.runtime_contractor("greedy", algebra=_NS()): + cons.set_contraction_algebra(_NS()) + try: + with cons.runtime_contractor("greedy"): + assert isinstance(cons.get_contraction_algebra(), _NS) + # runtime_contractor saves/restores algebra independently of the contractor, + # so algebra set before entering survives unchanged. assert isinstance(cons.get_contraction_algebra(), _NS) - assert cons.get_contraction_algebra() is prev + finally: + cons.set_contraction_algebra(prev) def test_set_function_contractor_restores_algebra(): prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(_NS()) + try: - @cons.set_function_contractor("greedy", algebra=_NS()) - def f(): - return cons.get_contraction_algebra() + @cons.set_function_contractor("greedy") + def f(): + return cons.get_contraction_algebra() - inside = f() - assert isinstance(inside, _NS) # algebra active during the call - assert cons.get_contraction_algebra() is prev # restored after + inside = f() + assert isinstance(inside, _NS) # algebra active during the call + # after the call, algebra survives unchanged (it was not set by the decorator) + assert isinstance(cons.get_contraction_algebra(), _NS) + finally: + cons.set_contraction_algebra(prev) # --- Task 3: cons.py guards (_merge_single_gates skip + _base routing) --- From 3ad432c5a2be5dfa50c77525234f9a06d8849fc6 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 16:05:12 +0800 Subject: [PATCH 06/16] refactor: replace prefer_einsum attr with get_contractor_kwargs() method - 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 --- applications/bcomplex32_algebra.py | 12 +- .../plans/2026-07-19-contractor-kwargs.md | 173 ++++++++++++++++++ .../2026-07-19-contractor-kwargs-design.md | 80 ++++++++ ...-prefer-einsum-contractor-kwargs-design.md | 83 +++++++++ examples/tropical_ising.py | 2 - pyproject.toml | 1 - tensorcircuit/cons.py | 2 +- tensorcircuit/contraction_algebra/base.py | 24 ++- tests/test_bcomplex32_algebra.py | 2 +- tests/test_contraction_algebra.py | 6 + tests/test_tropical.py | 7 +- 11 files changed, 370 insertions(+), 22 deletions(-) create mode 100644 docs/superpowers/plans/2026-07-19-contractor-kwargs.md create mode 100644 docs/superpowers/specs/2026-07-19-contractor-kwargs-design.md create mode 100644 docs/superpowers/specs/2026-07-19-prefer-einsum-contractor-kwargs-design.md diff --git a/applications/bcomplex32_algebra.py b/applications/bcomplex32_algebra.py index a04281ed..9206d168 100644 --- a/applications/bcomplex32_algebra.py +++ b/applications/bcomplex32_algebra.py @@ -33,10 +33,10 @@ def _complex_to_pair(be: Backend, t: Tensor) -> Tensor: def _pair_to_complex(be: Backend, pair: Tensor) -> Tensor: - """pair of bf16 -> complex64 tensor (recombine; no copy risk via cast).""" - re = be.cast(pair[..., 0], "float32") - im = be.cast(pair[..., 1], "float32") - return be.cast(re + 1j * im, "complex64") + """pair of bf16 -> complex tensor (recombine; no copy risk via cast).""" + re = be.cast(pair[..., 0], cons.rdtypestr) + im = be.cast(pair[..., 1], cons.rdtypestr) + return be.cast(re + 1j * im, cons.dtypestr) def _pair_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: @@ -175,7 +175,9 @@ def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor] class ComplexPairAlgebra(ContractionAlgebra): name = "bcomplex32_pair" representation = PairBf16Representation() - prefer_einsum = True # pair operands carry a trailing storage axis + + def get_contractor_kwargs(self) -> dict: + return {"prefer_einsum": True} def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: return _pair_tensordot(be, a, b, axes) diff --git a/docs/superpowers/plans/2026-07-19-contractor-kwargs.md b/docs/superpowers/plans/2026-07-19-contractor-kwargs.md new file mode 100644 index 00000000..5d010b01 --- /dev/null +++ b/docs/superpowers/plans/2026-07-19-contractor-kwargs.md @@ -0,0 +1,173 @@ +# Replace `prefer_einsum` with `get_contractor_kwargs()` — Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Replace the `prefer_einsum: bool` class attribute on `ContractionAlgebra` with a `get_contractor_kwargs()` method, keeping the same behavior while decoupling the ABC from cotengra internals. + +**Architecture:** A single new method on the ABC returns a dict of kwargs to forward to `ctg.core.make_contractor`. The default returns `{}`. `ComplexPairAlgebra` overrides to return `{"prefer_einsum": True}`. `cons.py` calls the method instead of reading the attribute. + +**Tech Stack:** Pure Python — no new dependencies. + +## Global Constraints + +- No behavior change for any algebra +- All existing tests must pass +- Three files touched: `base.py`, `bcomplex32_algebra.py`, `cons.py` + +--- + +### Task 1: Add `get_contractor_kwargs()` to ABC, remove `prefer_einsum` attribute + +**Files:** +- Modify: `tensorcircuit/contraction_algebra/base.py:59-69` + +**Interfaces:** +- Produces: `ContractionAlgebra.get_contractor_kwargs() -> dict` (default returns `{}`) + +- [ ] **Step 1: Replace `prefer_einsum` attribute with `get_contractor_kwargs()` method** + +Replace lines 59-69: + +```python + name: str = "abstract" + representation: Representation = IdentityRepresentation() + # When operands carry a trailing non-physical storage axis (e.g. the + # complex pair), cotengra's tensordot mode post-transposes results via + # autoray and mishandles that extra axis (ValueError: axes don't match array) + # -- set True to force einsum-only execution, which forwards operands verbatim + # to ``einsum`` and skips the autoray transpose. Default False keeps tensordot + # mode: tropical config-recovery backtracking depends on the tensordot + # intermediate layout and would break under forced einsum. + prefer_einsum: bool = False +``` + +With: + +```python + name: str = "abstract" + representation: Representation = IdentityRepresentation() + + def get_contractor_kwargs(self) -> dict: + """Extra kwargs forwarded to cotengra's ``make_contractor``. + + Override to return ``{'prefer_einsum': True}`` when your algebra's + ``tensordot`` kernel carries non-physical storage axes (e.g. the + complex pair axis) that cotengra's post-tensordot autoray + transpose would mishandle (ValueError: axes don't match array). + ``prefer_einsum=True`` forces einsum-only execution, which skips + the transpose entirely. + + Default ``{}`` keeps the standard tensordot+einsum mix — required + by tropical config-recovery backtracking, which depends on the + tensordot intermediate layout. + """ + return {} +``` + +- [ ] **Step 2: Verify the file is valid Python** + +Run: `D:/Software/miniconda3/envs/tcng/python.exe -c "from tensorcircuit.contraction_algebra.base import ContractionAlgebra; print(ContractionAlgebra().get_contractor_kwargs())"` +Expected: `{}` + +- [ ] **Step 3: Commit** + +```bash +git add tensorcircuit/contraction_algebra/base.py +git commit -m "refactor: replace prefer_einsum attr with get_contractor_kwargs() method on ContractionAlgebra" +``` + +--- + +### Task 2: Override `get_contractor_kwargs()` in `ComplexPairAlgebra` + +**Files:** +- Modify: `applications/bcomplex32_algebra.py:175-178` + +**Interfaces:** +- Consumes: `ContractionAlgebra.get_contractor_kwargs() -> dict` (from Task 1) + +- [ ] **Step 1: Replace `prefer_einsum` class attribute with method override** + +Replace line 178: + +```python + prefer_einsum = True # pair operands carry a trailing storage axis +``` + +With: + +```python + def get_contractor_kwargs(self) -> dict: + return {"prefer_einsum": True} +``` + +The class now looks like: + +```python +class ComplexPairAlgebra(ContractionAlgebra): + name = "bcomplex32_pair" + representation = PairBf16Representation() + + def get_contractor_kwargs(self) -> dict: + return {"prefer_einsum": True} + + def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: + ... +``` + +- [ ] **Step 2: Verify** + +Run: `D:/Software/miniconda3/envs/tcng/python.exe -c "from applications.bcomplex32_algebra import ComplexPairAlgebra; print(ComplexPairAlgebra().get_contractor_kwargs())"` +Expected: `{'prefer_einsum': True}` + +- [ ] **Step 3: Commit** + +```bash +git add applications/bcomplex32_algebra.py +git commit -m "refactor: use get_contractor_kwargs() override in ComplexPairAlgebra" +``` + +--- + +### Task 3: Update `cons.py` to call `get_contractor_kwargs()` + +**Files:** +- Modify: `tensorcircuit/cons.py:784-785` + +**Interfaces:** +- Consumes: `ContractionAlgebra.get_contractor_kwargs() -> dict` (from Task 1) + +- [ ] **Step 1: Replace `alg.prefer_einsum` with `**alg.get_contractor_kwargs()`** + +Replace lines 784-785: + +```python + contractor = ctg.core.make_contractor( + tree, implementation=impl, prefer_einsum=alg.prefer_einsum + ) +``` + +With: + +```python + contractor = ctg.core.make_contractor( + tree, implementation=impl, **alg.get_contractor_kwargs() + ) +``` + +- [ ] **Step 2: Verify no remaining references to `prefer_einsum` on algebra** + +Run: `cd "e:\Study\.AShare\OneDrive\OneDriveSync\session\tc\tensorcircuit-ng" && grep -rn "prefer_einsum" tensorcircuit/ applications/ --include="*.py" | grep -v test | grep -v __pycache__` +Expected: only `bcomplex32_algebra.py` (inside the new method returning the dict), and possibly `base.py` docstring mentioning it as an example + +- [ ] **Step 3: Run existing tests** + +Run: `D:/Software/miniconda3/envs/tcng/python.exe -m pytest tests/test_bcomplex32_algebra.py tests/test_contraction_algebra.py tests/test_tropical.py -v` +Expected: all pass + +- [ ] **Step 4: Commit** + +```bash +git add tensorcircuit/cons.py +git commit -m "refactor: consume alg.get_contractor_kwargs() in _algebraic_base_contraction" +``` \ No newline at end of file diff --git a/docs/superpowers/specs/2026-07-19-contractor-kwargs-design.md b/docs/superpowers/specs/2026-07-19-contractor-kwargs-design.md new file mode 100644 index 00000000..c01955d0 --- /dev/null +++ b/docs/superpowers/specs/2026-07-19-contractor-kwargs-design.md @@ -0,0 +1,80 @@ +# ContractionAlgebra: replace `prefer_einsum` class attribute with `get_contractor_kwargs()` method + +**Date:** 2026-07-19 +**Status:** draft +**PR:** feat/contraction-algebra-tropical + +## Motivation + +`ContractionAlgebra.prefer_einsum` is a cotengra-specific escape hatch that leaks +implementation details into the algebra ABC. The reviewer asked: "but you can +customize the tensordot, can you avoid this?" — the flag exists **not** because +`alg.tensordot` itself is broken, but because cotengra's `Contractor` applies a +post-`tensordot` autoray transpose whose `perm` doesn't know about non-physical +storage axes (e.g., the bf16 pair axis). The custom `tensordot` cannot prevent +this transpose, so `prefer_einsum=True` is the only way to skip it. + +However, `prefer_einsum` should not be a first-class ABC attribute. Instead, let +the algebra declare **contractor-level options** through a method whose return +value is transparently forwarded. + +## Design + +### ABC: `get_contractor_kwargs() -> dict` + +```python +class ContractionAlgebra(ABC): + def get_contractor_kwargs(self) -> dict: + """Extra kwargs forwarded to cotengra's make_contractor. + + Override to return ``{'prefer_einsum': True}`` when your kernel + carries non-physical storage axes that cotengra's post-tensordot + autoray transpose cannot handle. + """ + return {} +``` + +- Remove `prefer_einsum: bool = False`. +- Keep the explanatory comment, moved into the method docstring. + +### `ComplexPairAlgebra`: override + +```python +class ComplexPairAlgebra(ContractionAlgebra): + def get_contractor_kwargs(self) -> dict: + return {"prefer_einsum": True} +``` + +Remove the `prefer_einsum = True` class attribute. + +### `cons.py`: consume the method + +```python +# Before: +contractor = ctg.core.make_contractor( + tree, implementation=impl, prefer_einsum=alg.prefer_einsum +) + +# After: +contractor = ctg.core.make_contractor( + tree, implementation=impl, **alg.get_contractor_kwargs() +) +``` + +## Files changed + +| File | Change | +|------|--------| +| `tensorcircuit/contraction_algebra/base.py` | Remove `prefer_einsum` attr; add `get_contractor_kwargs()` method | +| `applications/bcomplex32_algebra.py` | Replace `prefer_einsum = True` with `get_contractor_kwargs()` override | +| `tensorcircuit/cons.py` | Call `alg.get_contractor_kwargs()` instead of `alg.prefer_einsum` | + +## Non-goals + +- Does not change cotengra's internal behavior +- Does not add new escape hatches — only re-packages the existing one +- No behavior change for `StandardAlgebra` or `TropicalAlgebra` + +## Reviewer reply + +> Good catch — `prefer_einsum` is now hidden behind a `get_contractor_kwargs()` method on the ABC so the algebra itself doesn't expose cotengra-specific flags as class attributes. \ No newline at end of file diff --git a/docs/superpowers/specs/2026-07-19-prefer-einsum-contractor-kwargs-design.md b/docs/superpowers/specs/2026-07-19-prefer-einsum-contractor-kwargs-design.md new file mode 100644 index 00000000..972f59be --- /dev/null +++ b/docs/superpowers/specs/2026-07-19-prefer-einsum-contractor-kwargs-design.md @@ -0,0 +1,83 @@ +# Extract `prefer_einsum` from ContractionAlgebra ABC into `get_contractor_kwargs()` + +Date: 2026-07-19 + +## Motivation + +`ContractionAlgebra` currently carries a `prefer_einsum: bool` class attribute whose +sole purpose is to be forwarded to `ctg.core.make_contractor(tree, prefer_einsum=...)`. +It exists because cotengra's tensordot mode applies a post-contraction transpose whose +permutation axes do not account for non-physical storage axes (e.g. the trailing +complex pair axis), causing `ValueError: axes don't match array`. + +This flag is an implementation detail of *how* cotengra executes contractions and +does not belong in the algebra ABC. The reviewer correctly observed that since the +algebra already customizes `tensordot`, a second escape-hatch flag seems redundant +— but the actual problem is in cotengra's internal transpose, not in our tensordot. + +## Design + +Replace the hard attribute with a method: + +```python +class ContractionAlgebra(ABC): + def get_contractor_kwargs(self) -> Dict[str, Any]: + """Extra keyword arguments forwarded to ``ctg.core.make_contractor``. + + Override to return ``{"prefer_einsum": True}`` when your contraction + kernels produce tensors with non-physical storage axes (e.g. the + complex pair axis) that cotengra's post-tensordot transpose + cannot handle. + + The default empty dict preserves standard tensordot mode, which + tropical config-recovery backtracking depends on. + """ + return {} +``` + +`ComplexPairAlgebra` overrides it: + +```python +class ComplexPairAlgebra(ContractionAlgebra): + def get_contractor_kwargs(self) -> Dict[str, Any]: + return {"prefer_einsum": True} +``` + +In `cons.py` the dispatch changes from: + +```python +contractor = ctg.core.make_contractor( + tree, implementation=impl, prefer_einsum=alg.prefer_einsum +) +``` + +to: + +```python +contractor = ctg.core.make_contractor( + tree, implementation=impl, **alg.get_contractor_kwargs() +) +``` + +## Trade-offs + +- **Pro:** `ContractionAlgebra` ABC no longer exposes cotengra's `prefer_einsum` + parameter as a first-class attribute. +- **Pro:** The dict return type is extensible — future cotengra kwargs + (e.g. `strip_exponent`) can be supplied without changing the ABC again. +- **Pro:** Default `{}` keeps standard algebras unchanged; tropical contracts + through the same path as before. +- **Con:** The method signature still implies that the algebra knows about + cotengra's constructor interface. This coupling is acceptable because the + algebra's *only* consumer is `_algebraic_base_contraction`, which builds a + cotengra contractor. + +## Affected files + +| File | Change | +|------|--------| +| `tensorcircuit/contraction_algebra/base.py` | Replace `prefer_einsum: bool` attribute with `get_contractor_kwargs()` method | +| `applications/bcomplex32_algebra.py` | Replace `prefer_einsum = True` with `get_contractor_kwargs()` override | +| `tensorcircuit/cons.py` | Call `alg.get_contractor_kwargs()` instead of reading `alg.prefer_einsum` | + +Tests do not need changes — the runtime behavior is identical. \ No newline at end of file diff --git a/examples/tropical_ising.py b/examples/tropical_ising.py index 9dbd1abb..0e8124c3 100644 --- a/examples/tropical_ising.py +++ b/examples/tropical_ising.py @@ -1,6 +1,4 @@ """Tropical (max-plus) Ising ground-state energy via TensorCircuit-NG contraction. - -Run: PYTHONPATH= python examples/tropical_ising.py """ import itertools diff --git a/pyproject.toml b/pyproject.toml index 0a2a7db0..c1984d1b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,6 @@ filterwarnings = [ "ignore::DeprecationWarning", "ignore:Explicitly requested dtype*:UserWarning", ] -pythonpath = ["."] [tool.mypy] diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 4f4d6af5..4fc1c239 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -782,7 +782,7 @@ def _run_contraction( functools.partial(alg.tensordot, kbe), ) contractor = ctg.core.make_contractor( - tree, implementation=impl, prefer_einsum=alg.prefer_einsum + tree, implementation=impl, **alg.get_contractor_kwargs() ) final = contractor(*raw_tensors) elif not strip_exponent: diff --git a/tensorcircuit/contraction_algebra/base.py b/tensorcircuit/contraction_algebra/base.py index 4e680c14..6293ec6c 100644 --- a/tensorcircuit/contraction_algebra/base.py +++ b/tensorcircuit/contraction_algebra/base.py @@ -59,14 +59,22 @@ class ContractionAlgebra(ABC): name: str = "abstract" representation: Representation = IdentityRepresentation() - # When operands carry a trailing non-physical storage axis (e.g. the - # complex pair), cotengra's tensordot mode post-transposes results via - # autoray and mishandles that extra axis (ValueError: axes don't match array) - # -- set True to force einsum-only execution, which forwards operands verbatim - # to ``einsum`` and skips the autoray transpose. Default False keeps tensordot - # mode: tropical config-recovery backtracking depends on the tensordot - # intermediate layout and would break under forced einsum. - prefer_einsum: bool = False + + def get_contractor_kwargs(self) -> dict: + """Extra kwargs forwarded to cotengra's ``make_contractor``. + + Override to return ``{'prefer_einsum': True}`` when your algebra's + ``tensordot`` kernel carries non-physical storage axes (e.g. the + complex pair axis) that cotengra's post-tensordot autoray + transpose would mishandle (ValueError: axes don't match array). + ``prefer_einsum=True`` forces einsum-only execution, which skips + the transpose entirely. + + Default ``{}`` keeps the standard tensordot+einsum mix — required + by tropical config-recovery backtracking, which depends on the + tensordot intermediate layout. + """ + return {} @abstractmethod def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: diff --git a/tests/test_bcomplex32_algebra.py b/tests/test_bcomplex32_algebra.py index d182d9de..b9d8a50a 100644 --- a/tests/test_bcomplex32_algebra.py +++ b/tests/test_bcomplex32_algebra.py @@ -111,7 +111,7 @@ def test_pair_einsum_keeps_bfloat16_dtype(): def test_bf16_ghz8_runs_and_matches_native(): """T3: the 8-qubit GHZ that previously crashed (cotengra autoray transpose on a pair result) now runs under bcomplex32 and matches native within bf16 - tolerance. prefer_einsum=True avoids the transpose path; the genuine-bf16 + tolerance. get_contractor_kwargs returns prefer_einsum=True, which avoids the transpose path; the genuine-bf16 kernel keeps intermediates bf16 end-to-end. """ diff --git a/tests/test_contraction_algebra.py b/tests/test_contraction_algebra.py index 068a8d67..5e3ceeea 100644 --- a/tests/test_contraction_algebra.py +++ b/tests/test_contraction_algebra.py @@ -402,3 +402,9 @@ def einsum(self, be, eq, *ops): assert "count" in cons._aux_outputs() # aux stashed without crashing finally: cons.set_contraction_algebra(prev) + + +def test_get_contractor_kwargs_default(): + assert StandardAlgebra().get_contractor_kwargs() == {} + # Concrete algebra without override inherits the default + assert _NS().get_contractor_kwargs() == {} diff --git a/tests/test_tropical.py b/tests/test_tropical.py index be2af826..33a6b681 100644 --- a/tests/test_tropical.py +++ b/tests/test_tropical.py @@ -20,7 +20,6 @@ MaxPlusAlgebra, MaxPlusTrackingAlgebra, CountingTropicalAlgebra, - CountingRepresentation, tropical, counting_tropical, recover_configuration, @@ -539,10 +538,10 @@ def test_preprocessing_does_not_corrupt_tropical(): def test_tropical_public_api_surface(): + assert MaxPlusAlgebra is not None + assert MaxPlusTrackingAlgebra is not None + assert CountingTropicalAlgebra is not None for name in [ - "MaxPlusAlgebra", - "MaxPlusTrackingAlgebra", - "CountingTropicalAlgebra", "tropical", "counting_tropical", "recover_configuration", From df2776b893d2216915ce88fa9e5359967002fd79 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 16:19:07 +0800 Subject: [PATCH 07/16] refactor: flatten contraction_algebra subpackage into single file MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- tensorcircuit/cons.py | 4 +-- .../base.py => contraction_algebra.py} | 26 +++++++++++------- tensorcircuit/contraction_algebra/__init__.py | 27 ------------------- tests/test_contraction_algebra.py | 2 +- 4 files changed, 20 insertions(+), 39 deletions(-) rename tensorcircuit/{contraction_algebra/base.py => contraction_algebra.py} (83%) delete mode 100644 tensorcircuit/contraction_algebra/__init__.py diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 4fc1c239..5f5d3c82 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -140,9 +140,9 @@ def set_tensornetwork_backend( set_tensornetwork_backend() # --- ContractionAlgebra state (mirrors dtypestr/set_dtype pattern) --- -from .contraction_algebra import StandardAlgebra as _StandardAlgebra -from .contraction_algebra.base import ( +from .contraction_algebra import ( ContractionAlgebra as _ContractionAlgebra, + StandardAlgebra as _StandardAlgebra, ) _contraction_algebra: _ContractionAlgebra = _StandardAlgebra() diff --git a/tensorcircuit/contraction_algebra/base.py b/tensorcircuit/contraction_algebra.py similarity index 83% rename from tensorcircuit/contraction_algebra/base.py rename to tensorcircuit/contraction_algebra.py index 6293ec6c..aad34704 100644 --- a/tensorcircuit/contraction_algebra/base.py +++ b/tensorcircuit/contraction_algebra.py @@ -1,15 +1,15 @@ -"""ContractionAlgebra: generic interface for swapping contraction primitives + +"""ContractionAlgebra: a generic interface for swapping contraction primitives + boundary representation, consulted by ``cons._algebraic_base_contraction``. -Two ABCs: -- ``ContractionAlgebra``: the arithmetic (tensordot/einsum kernels) + observation - hooks. Carries a ``Representation`` (boundary codec), default identity. -- ``Representation``: boundary encode (leaves -> physical storage) / decode - (final -> primary tensor + aux). Covers storage-split (complex<->pair) and - precision cast (f64<->bf16). Default identity. +Implement ``ContractionAlgebra`` (with a ``Representation``) and activate via +``cons.set_contraction_algebra(...)`` -- the single entry-point for switching +contraction algebras. The in-source ``cons._base`` +routes any non-standard algebra to ``_algebraic_base_contraction``, which runs +encode -> algebra kernels -> decode; no monkey-patching is required. -kernel must be closed over whatever its Representation produces (author's -contract; Representation is bundled in the algebra so users cannot mis-pair). +Reference applications: ``applications/tropical_algebra.py`` +(max-plus / counting / tracking) and ``applications/bcomplex32_algebra.py`` +(bf16 pair). """ from abc import ABC, abstractmethod @@ -101,3 +101,11 @@ def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: return be.einsum(eq, *operands) + + +__all__ = [ + "ContractionAlgebra", + "StandardAlgebra", + "Representation", + "IdentityRepresentation", +] \ No newline at end of file diff --git a/tensorcircuit/contraction_algebra/__init__.py b/tensorcircuit/contraction_algebra/__init__.py deleted file mode 100644 index b4c76cda..00000000 --- a/tensorcircuit/contraction_algebra/__init__.py +++ /dev/null @@ -1,27 +0,0 @@ -"""ContractionAlgebra: a generic interface for swapping contraction primitives + -boundary representation, consulted by ``cons._algebraic_base_contraction``. - -Implement ``ContractionAlgebra`` (with a ``Representation``) and activate via -``cons.set_contraction_algebra(...)`` -- the single entry-point for switching -contraction algebras. The in-source ``cons._base`` -routes any non-standard algebra to ``_algebraic_base_contraction``, which runs -encode -> algebra kernels -> decode; no monkey-patching is required. - -Reference applications: ``applications/tropical_algebra.py`` -(max-plus / counting / tracking) and ``applications/bcomplex32_algebra.py`` -(bf16 pair). -""" - -from .base import ( - ContractionAlgebra, - StandardAlgebra, - Representation, - IdentityRepresentation, -) - -__all__ = [ - "ContractionAlgebra", - "StandardAlgebra", - "Representation", - "IdentityRepresentation", -] diff --git a/tests/test_contraction_algebra.py b/tests/test_contraction_algebra.py index 5e3ceeea..22e6ef30 100644 --- a/tests/test_contraction_algebra.py +++ b/tests/test_contraction_algebra.py @@ -1,7 +1,7 @@ import numpy as np import tensorcircuit as tc import tensorcircuit.cons as cons -from tensorcircuit.contraction_algebra.base import ContractionAlgebra, StandardAlgebra +from tensorcircuit.contraction_algebra import ContractionAlgebra, StandardAlgebra def test_standard_tensordot_matches_backend(): From 3c78e81448e05a0319d209ee4322ef073b168571 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 16:47:36 +0800 Subject: [PATCH 08/16] fix: strict mypy Dict[str, Any] annotation + black formatting --- applications/bcomplex32_algebra.py | 2 +- examples/tropical_ising.py | 3 +-- tensorcircuit/contraction_algebra.py | 4 ++-- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/applications/bcomplex32_algebra.py b/applications/bcomplex32_algebra.py index 9206d168..73dcf2ab 100644 --- a/applications/bcomplex32_algebra.py +++ b/applications/bcomplex32_algebra.py @@ -176,7 +176,7 @@ class ComplexPairAlgebra(ContractionAlgebra): name = "bcomplex32_pair" representation = PairBf16Representation() - def get_contractor_kwargs(self) -> dict: + def get_contractor_kwargs(self) -> Dict[str, Any]: return {"prefer_einsum": True} def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: diff --git a/examples/tropical_ising.py b/examples/tropical_ising.py index 0e8124c3..7324cad8 100644 --- a/examples/tropical_ising.py +++ b/examples/tropical_ising.py @@ -1,5 +1,4 @@ -"""Tropical (max-plus) Ising ground-state energy via TensorCircuit-NG contraction. -""" +"""Tropical (max-plus) Ising ground-state energy via TensorCircuit-NG contraction.""" import itertools import numpy as np diff --git a/tensorcircuit/contraction_algebra.py b/tensorcircuit/contraction_algebra.py index aad34704..4e7cfc79 100644 --- a/tensorcircuit/contraction_algebra.py +++ b/tensorcircuit/contraction_algebra.py @@ -60,7 +60,7 @@ class ContractionAlgebra(ABC): name: str = "abstract" representation: Representation = IdentityRepresentation() - def get_contractor_kwargs(self) -> dict: + def get_contractor_kwargs(self) -> Dict[str, Any]: """Extra kwargs forwarded to cotengra's ``make_contractor``. Override to return ``{'prefer_einsum': True}`` when your algebra's @@ -108,4 +108,4 @@ def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: "StandardAlgebra", "Representation", "IdentityRepresentation", -] \ No newline at end of file +] From 2b161004cf7be05d47e2dd3423c3542bf275ba02 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 16:57:55 +0800 Subject: [PATCH 09/16] refactor: move contraction_algebra import to top of cons.py --- tensorcircuit/cons.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 5f5d3c82..0dec8695 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -22,6 +22,10 @@ from .backends.numpy_backend import NumpyBackend from .backends import get_backend from .simplify import _multi_remove +from .contraction_algebra import ( + ContractionAlgebra as _ContractionAlgebra, + StandardAlgebra as _StandardAlgebra, +) logger = logging.getLogger(__name__) @@ -140,10 +144,6 @@ def set_tensornetwork_backend( set_tensornetwork_backend() # --- ContractionAlgebra state (mirrors dtypestr/set_dtype pattern) --- -from .contraction_algebra import ( - ContractionAlgebra as _ContractionAlgebra, - StandardAlgebra as _StandardAlgebra, -) _contraction_algebra: _ContractionAlgebra = _StandardAlgebra() From 986588b7c8365c4734e9fc19bdc29a329852a5cf Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 18:19:17 +0800 Subject: [PATCH 10/16] fix: address PR review minor issues (code quality + docs) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- CHANGELOG.md | 2 +- applications/tropical_algebra.py | 25 ++++++------------------- requirements/requirements-dev.txt | 2 +- tensorcircuit/cons.py | 11 ++++++++++- tests/test_bcomplex32_algebra.py | 4 +++- tests/test_contraction_algebra.py | 8 ++++---- 6 files changed, 25 insertions(+), 27 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b5552469..f1494c42 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### Added -- add unified ContractionAlgebra interface (`set_contractor(algebra=...)` + boundary +- add unified ContractionAlgebra interface (`set_contraction_algebra(...)` + boundary encode/decode) with tropical (max-plus) and complex pair-algebra reference applications. diff --git a/applications/tropical_algebra.py b/applications/tropical_algebra.py index a2918385..d6690e0d 100644 --- a/applications/tropical_algebra.py +++ b/applications/tropical_algebra.py @@ -229,6 +229,9 @@ def _resolve_repeats( """ if len(set(idxs)) != len(idxs): resolved = "".join(dict.fromkeys(idxs)) + # Diagonal extraction: exactly one element contributes per output position, + # so the standard (sum, multiply) einsum is equivalent to max — both + # reduce to identity for a singleton set of values. e = be.einsum("".join(idxs) + "->" + resolved, pair[..., 0]) n = be.einsum("".join(idxs) + "->" + resolved, pair[..., 1]) return e, n, list(resolved) @@ -484,7 +487,9 @@ def _tracking_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: contract_labels = [c for c in all_idx if c not in out_labels] # Full-layout pair sum (same construction as _tropical_einsum). - s = _expand_to_layout_pair(be, a, b, ia, ib, all_idx) + s = _expand_to_layout(be, a, ia, all_idx) + _expand_to_layout( + be, b, ib, all_idx + ) if not contract_labels: # Pure outer product (hyperedge with no contraction): no argmax to take. @@ -522,24 +527,6 @@ def _tracking_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: return _tropical_einsum(be, eq, *operands) -def _expand_to_layout_pair( - be: Backend, - a: Tensor, - b: Tensor, - ia: Sequence[str], - ib: Sequence[str], - all_idx: Sequence[str], -) -> Tensor: - """Build the full-layout pair-sum ``a + b`` broadcast over ``all_idx``. - - Mirrors the pair-sum construction inside ``_tropical_einsum`` (factors - ``_expand_to_layout`` for both operands). Kept local so this module does not - reach into a private helper whose signature may change. - """ - - return _expand_to_layout(be, a, ia, all_idx) + _expand_to_layout(be, b, ib, all_idx) - - class MaxPlusTrackingAlgebra(MaxPlusAlgebra): """Max-plus algebra that records the per-step argmax for config recovery. diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt index 1b6edaff..a4d5d865 100644 --- a/requirements/requirements-dev.txt +++ b/requirements/requirements-dev.txt @@ -3,7 +3,7 @@ pytest==7.4.4 pytest-cov pytest-benchmark pytest-xdist -black[jupyter]==26.5.1 +black[jupyter] sphinx>=4.0 pytest-lazy-fixture pylint==3.2.6 diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 0dec8695..75f1da3c 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -162,6 +162,12 @@ def set_contraction_algebra(alg: _ContractionAlgebra) -> None: # primary tensor is returned through the normal contraction return value, while # aux (e.g. counting degeneracy, sharing the primary's physical axis order) is # stashed here for the caller to read via ``_aux_outputs``. +# +# Thread-safety: this global, along with ``_trace`` and ``_ctx`` in +# ``applications/tropical_algebra.py``, is module-level mutable state cleared +# per contraction. Concurrent contractions will corrupt each other. Acceptable +# for the single-threaded research use case; see also +# ``docs/superpowers/plans/2026-07-19-L3-bf16-backend-generic.md``. _aux_outputs_store: Dict[str, Any] = {} @@ -838,6 +844,7 @@ def _algebraic_base_contraction( rep = alg.representation ns = not isinstance(alg, _StandardAlgebra) kbe = backend if ns else be # algebra kernels need the tc backend (max/argmax/...) + # be and backend are normally the same object: tn.set_default_backend syncs them. # Unconditional clear: every contraction (standard or non-standard) wipes the # aux side-channel so a subsequent ``degeneracy()`` cannot read a stale count @@ -1036,7 +1043,7 @@ def _base( _ns_alg = not isinstance(_contraction_algebra, _StandardAlgebra) if use_primitives is True or _ns_alg or (use_primitives is None and has_hyperedges): # ========================================== - # NEW ALGEBRAIC EXECUTION PATH (Opt-in) + # ALGEBRAIC EXECUTION PATH # ========================================== return _algebraic_base_contraction( nodes, algorithm, output_edge_order, ignore_edge_order, **kws @@ -1046,6 +1053,8 @@ def _base( # ORIGINAL EXECUTION PATH (100% Backward Compatible) # ========================================== + _stash_aux_outputs({}) # clear stale aux from prior non-standard contraction + for edge in edges: if not edge.is_disabled: # if its disabled we already contracted it if edge.is_trace(): diff --git a/tests/test_bcomplex32_algebra.py b/tests/test_bcomplex32_algebra.py index b9d8a50a..d002cb09 100644 --- a/tests/test_bcomplex32_algebra.py +++ b/tests/test_bcomplex32_algebra.py @@ -126,4 +126,6 @@ def ghz(n): with bcomplex32(): got = ghz(8) assert got.shape == ref.shape - assert np.allclose(got, ref, rtol=5e-2), f"max abs diff = {np.abs(got - ref).max()}" + assert np.allclose( + got, ref, rtol=1.5e-2 + ), f"max abs diff = {np.abs(got - ref).max()}" diff --git a/tests/test_contraction_algebra.py b/tests/test_contraction_algebra.py index 22e6ef30..0f331a2e 100644 --- a/tests/test_contraction_algebra.py +++ b/tests/test_contraction_algebra.py @@ -279,6 +279,7 @@ def f(): def test_merge_single_gates_skipped_under_nonstandard_algebra(monkeypatch): + import numpy as np import tensornetwork as tn prev = cons.get_contraction_algebra() @@ -291,10 +292,9 @@ def boom(*a, **k): ) monkeypatch.setattr(tn, "contract_parallel", boom) - out = cons._merge_single_gates( - ["fake_node"], 7 - ) # guard fires before any node access - assert out == (["fake_node"], 7) + node = tn.Node(np.zeros(2)) + out = cons._merge_single_gates([node], 7) + assert out[0] == [node] and out[1] == 7 finally: cons.set_contraction_algebra(prev) From c8e8aea4bc76e18821e16a6a258a73bd6303a5f5 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 19:07:21 +0800 Subject: [PATCH 11/16] refactor: default _contraction_algebra to None instead of StandardAlgebra - _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 --- tensorcircuit/cons.py | 15 ++++++++------- tests/test_bcomplex32_algebra.py | 7 +++---- tests/test_contraction_algebra.py | 2 +- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 75f1da3c..82811875 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -145,14 +145,15 @@ def set_tensornetwork_backend( # --- ContractionAlgebra state (mirrors dtypestr/set_dtype pattern) --- -_contraction_algebra: _ContractionAlgebra = _StandardAlgebra() +_contraction_algebra: Optional[_ContractionAlgebra] = None +_standard = _StandardAlgebra() # singleton fallback for primitives/hyperedge path -def get_contraction_algebra() -> _ContractionAlgebra: +def get_contraction_algebra() -> Optional[_ContractionAlgebra]: return _contraction_algebra -def set_contraction_algebra(alg: _ContractionAlgebra) -> None: +def set_contraction_algebra(alg: Optional[_ContractionAlgebra]) -> None: global _contraction_algebra _contraction_algebra = alg @@ -341,7 +342,7 @@ def _sizen(node: tn.Node, is_log: bool = False) -> int: def _merge_single_gates( nodes: List[Any], total_size: Optional[int] = None ) -> Tuple[List[Any], int]: - if not isinstance(_contraction_algebra, _StandardAlgebra): + if _contraction_algebra is not None: # _merge_single_gates contracts via tn.contract_parallel (native sum-product), # bypassing the algebra kernel — skip it under a non-standard algebra. if total_size is None: @@ -840,9 +841,9 @@ def _algebraic_base_contraction( raw_tensors, input_sets, output_set, size_dict = _extract_topology(nodes) be = nodes[0].backend # tn backend: standard native ops + tn.Node wrap - alg = get_contraction_algebra() + alg = get_contraction_algebra() or _standard # fall back to standard when None rep = alg.representation - ns = not isinstance(alg, _StandardAlgebra) + ns = alg is not _standard kbe = backend if ns else be # algebra kernels need the tc backend (max/argmax/...) # be and backend are normally the same object: tn.set_default_backend syncs them. @@ -1040,7 +1041,7 @@ def _base( # 1. Resolve topology and check for hyperedges has_hyperedges = any(isinstance(n, tn.CopyNode) for n in nodes) - _ns_alg = not isinstance(_contraction_algebra, _StandardAlgebra) + _ns_alg = _contraction_algebra is not None if use_primitives is True or _ns_alg or (use_primitives is None and has_hyperedges): # ========================================== # ALGEBRAIC EXECUTION PATH diff --git a/tests/test_bcomplex32_algebra.py b/tests/test_bcomplex32_algebra.py index d002cb09..be9fa3e6 100644 --- a/tests/test_bcomplex32_algebra.py +++ b/tests/test_bcomplex32_algebra.py @@ -79,11 +79,10 @@ def test_bf16_wall_avoidance_canary(): st = np.asarray(c.state()) assert st.shape == (16,) # ran cleanly, no axis==edge crash import tensorcircuit.cons as cons - from tensorcircuit.contraction_algebra import StandardAlgebra - assert isinstance( - cons.get_contraction_algebra(), StandardAlgebra - ) # CM restored algebra, no leak + assert ( + cons.get_contraction_algebra() is None + ) # CM restored algebra to default (None), no leak c2 = tc.Circuit(2) c2.H(0) c2.cnot(0, 1) # subsequent native contraction diff --git a/tests/test_contraction_algebra.py b/tests/test_contraction_algebra.py index 0f331a2e..04dd5d8d 100644 --- a/tests/test_contraction_algebra.py +++ b/tests/test_contraction_algebra.py @@ -345,7 +345,7 @@ def test_standard_algebra_contraction_matches_native(): b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) tn.connect(a[1], b[0]) prev = cons.get_contraction_algebra() - assert isinstance(prev, StandardAlgebra) + assert prev is None # default is None, not StandardAlgebra cons.set_contraction_algebra(StandardAlgebra()) try: import opt_einsum From d69d7f79cb11c23f1f670fae4c2a26e926469042 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 21:24:09 +0800 Subject: [PATCH 12/16] refactor: default _contraction_algebra to None, clean up guard checks --- tensorcircuit/cons.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 82811875..510add89 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -146,14 +146,13 @@ def set_tensornetwork_backend( # --- ContractionAlgebra state (mirrors dtypestr/set_dtype pattern) --- _contraction_algebra: Optional[_ContractionAlgebra] = None -_standard = _StandardAlgebra() # singleton fallback for primitives/hyperedge path def get_contraction_algebra() -> Optional[_ContractionAlgebra]: return _contraction_algebra -def set_contraction_algebra(alg: Optional[_ContractionAlgebra]) -> None: +def set_contraction_algebra(alg: _ContractionAlgebra) -> None: global _contraction_algebra _contraction_algebra = alg @@ -841,9 +840,9 @@ def _algebraic_base_contraction( raw_tensors, input_sets, output_set, size_dict = _extract_topology(nodes) be = nodes[0].backend # tn backend: standard native ops + tn.Node wrap - alg = get_contraction_algebra() or _standard # fall back to standard when None + alg = get_contraction_algebra() rep = alg.representation - ns = alg is not _standard + ns = not isinstance(alg, _StandardAlgebra) kbe = backend if ns else be # algebra kernels need the tc backend (max/argmax/...) # be and backend are normally the same object: tn.set_default_backend syncs them. @@ -1041,11 +1040,11 @@ def _base( # 1. Resolve topology and check for hyperedges has_hyperedges = any(isinstance(n, tn.CopyNode) for n in nodes) - _ns_alg = _contraction_algebra is not None - if use_primitives is True or _ns_alg or (use_primitives is None and has_hyperedges): - # ========================================== - # ALGEBRAIC EXECUTION PATH - # ========================================== + if _contraction_algebra is not None: + return _algebraic_base_contraction( + nodes, algorithm, output_edge_order, ignore_edge_order, **kws + ) + if use_primitives is True or (use_primitives is None and has_hyperedges): return _algebraic_base_contraction( nodes, algorithm, output_edge_order, ignore_edge_order, **kws ) From dcb949070b268dd08c3a8d9366d55b7607195e3f Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 21:28:44 +0800 Subject: [PATCH 13/16] =?UTF-8?q?refactor:=20simplify=20=5Falgebraic=5Fbas?= =?UTF-8?q?e=5Fcontraction=20=E2=80=94=20ns=20=3D=20alg=20is=20not=20None?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tensorcircuit/cons.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 510add89..95c09d18 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -24,7 +24,7 @@ from .simplify import _multi_remove from .contraction_algebra import ( ContractionAlgebra as _ContractionAlgebra, - StandardAlgebra as _StandardAlgebra, + IdentityRepresentation, ) logger = logging.getLogger(__name__) @@ -841,14 +841,16 @@ def _algebraic_base_contraction( be = nodes[0].backend # tn backend: standard native ops + tn.Node wrap alg = get_contraction_algebra() - rep = alg.representation - ns = not isinstance(alg, _StandardAlgebra) - kbe = backend if ns else be # algebra kernels need the tc backend (max/argmax/...) + ns = alg is not None + + if ns: + rep = alg.representation + kbe = backend # algebra kernels need the tc backend (max/argmax/...) + else: + rep = IdentityRepresentation() # no-op; _decode_aux skips it when ns=False + kbe = be # be and backend are normally the same object: tn.set_default_backend syncs them. - # Unconditional clear: every contraction (standard or non-standard) wipes the - # aux side-channel so a subsequent ``degeneracy()`` cannot read a stale count - # left by an earlier counting contraction. Non-standard decodes then refill it. _stash_aux_outputs({}) if ns: if kws.get("strip_exponent", False): From a34521517c60167aa4a846d061263eb1e2b57587 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 21:30:45 +0800 Subject: [PATCH 14/16] test: update assertions for _contraction_algebra default None --- tests/test_bcomplex32_algebra.py | 4 +--- tests/test_contraction_algebra.py | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/test_bcomplex32_algebra.py b/tests/test_bcomplex32_algebra.py index be9fa3e6..ad0e56d6 100644 --- a/tests/test_bcomplex32_algebra.py +++ b/tests/test_bcomplex32_algebra.py @@ -80,9 +80,7 @@ def test_bf16_wall_avoidance_canary(): assert st.shape == (16,) # ran cleanly, no axis==edge crash import tensorcircuit.cons as cons - assert ( - cons.get_contraction_algebra() is None - ) # CM restored algebra to default (None), no leak + assert cons.get_contraction_algebra() is None # CM restored, no leak c2 = tc.Circuit(2) c2.H(0) c2.cnot(0, 1) # subsequent native contraction diff --git a/tests/test_contraction_algebra.py b/tests/test_contraction_algebra.py index 04dd5d8d..efe64467 100644 --- a/tests/test_contraction_algebra.py +++ b/tests/test_contraction_algebra.py @@ -345,7 +345,7 @@ def test_standard_algebra_contraction_matches_native(): b = tn.Node(rng.standard_normal((3, 4)).astype(np.complex64)) tn.connect(a[1], b[0]) prev = cons.get_contraction_algebra() - assert prev is None # default is None, not StandardAlgebra + assert prev is None cons.set_contraction_algebra(StandardAlgebra()) try: import opt_einsum From 01847765ad614d2d9b726f2e26401948f145edf4 Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Sun, 19 Jul 2026 21:46:56 +0800 Subject: [PATCH 15/16] fix: inline alg is not None for mypy type narrowing --- tensorcircuit/cons.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 95c09d18..6a2ac07d 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -152,7 +152,7 @@ def get_contraction_algebra() -> Optional[_ContractionAlgebra]: return _contraction_algebra -def set_contraction_algebra(alg: _ContractionAlgebra) -> None: +def set_contraction_algebra(alg: Optional[_ContractionAlgebra]) -> None: global _contraction_algebra _contraction_algebra = alg @@ -841,9 +841,8 @@ def _algebraic_base_contraction( be = nodes[0].backend # tn backend: standard native ops + tn.Node wrap alg = get_contraction_algebra() - ns = alg is not None - if ns: + if alg is not None: rep = alg.representation kbe = backend # algebra kernels need the tc backend (max/argmax/...) else: @@ -852,7 +851,7 @@ def _algebraic_base_contraction( # be and backend are normally the same object: tn.set_default_backend syncs them. _stash_aux_outputs({}) - if ns: + if alg is not None: if kws.get("strip_exponent", False): raise ValueError( "strip_exponent is incompatible with a non-standard ContractionAlgebra" @@ -866,7 +865,7 @@ def _algebraic_base_contraction( output_set, size_dict, algorithm, - ns, + alg is not None, alg, kbe, be, @@ -874,7 +873,7 @@ def _algebraic_base_contraction( ctg, ) - final, aux = _decode_aux(ns, rep, kbe, final, output_set) + final, aux = _decode_aux(alg is not None, rep, kbe, final, output_set) final_node = tn.Node(final, backend=be) @@ -890,10 +889,10 @@ def _algebraic_base_contraction( final_node.reorder_edges(list(output_edge_order)) # Apply the same output_edge_order permutation to aux (count co-indexed with energy) - if ns and aux: + if alg is not None and aux: _stash_permuted_aux(aux, output_edge_order, dangling_edges, kbe) - if kws.get("strip_exponent", False) and not ns: + if kws.get("strip_exponent", False) and alg is None: return final_node, exponent return final_node From bc7671c68119874e2d6f51f9531f1198ba7d5b9c Mon Sep 17 00:00:00 2001 From: DeanTMaxim <824411625@qq.com> Date: Mon, 20 Jul 2026 19:32:01 +0800 Subject: [PATCH 16/16] introduce pairtensor wrapper and apply contraction algebra review feedback --- applications/bcomplex32_algebra.py | 95 ++++------ applications/tropical_algebra.py | 128 +++++-------- .../plans/2026-07-19-contractor-kwargs.md | 173 ------------------ .../2026-07-19-contractor-kwargs-design.md | 80 -------- ...-prefer-einsum-contractor-kwargs-design.md | 83 --------- tensorcircuit/cons.py | 140 +++++++------- tensorcircuit/contraction_algebra.py | 45 +++++ tests/test_bcomplex32_algebra.py | 7 +- tests/test_contraction_algebra.py | 76 +++++++- 9 files changed, 265 insertions(+), 562 deletions(-) delete mode 100644 docs/superpowers/plans/2026-07-19-contractor-kwargs.md delete mode 100644 docs/superpowers/specs/2026-07-19-contractor-kwargs-design.md delete mode 100644 docs/superpowers/specs/2026-07-19-prefer-einsum-contractor-kwargs-design.md diff --git a/applications/bcomplex32_algebra.py b/applications/bcomplex32_algebra.py index 73dcf2ab..96c04356 100644 --- a/applications/bcomplex32_algebra.py +++ b/applications/bcomplex32_algebra.py @@ -1,9 +1,8 @@ """complex pair-algebra — a reference APPLICATION of ContractionAlgebra. -Pair repr: complex tensor = stack([re, im], axis=-1) of bf16. Contraction = 4 real +Pair repr: complex tensor split into PairTensor(re, im) of bf16. Contraction = 4 real bf16 matmuls (4M). Activated via ``cons.set_contraction_algebra(ComplexPairAlgebra())`` or -the ``bcomplex32()`` CM. encode/decode at the ``_algebraic_base_contraction`` boundary -keep the pair axis off tn.Node (dodges the axis==edge wall). +the ``bcomplex32()`` CM. PairTensor keeps the pair axis off tn.Node (dodges the axis==edge wall). """ from typing import Any, Dict, Iterator, List, Tuple @@ -12,7 +11,11 @@ import numpy as np import tensorcircuit.cons as cons -from tensorcircuit.contraction_algebra import ContractionAlgebra, Representation +from tensorcircuit.contraction_algebra import ( + ContractionAlgebra, + PairTensor, + Representation, +) Tensor = Any Backend = Any @@ -24,30 +27,39 @@ def _bf16_dtype() -> Any: return ml_dtypes.bfloat16 -def _complex_to_pair(be: Backend, t: Tensor) -> Tensor: - """complex tensor -> stack([re, im], axis=-1) of bf16.""" +def _complex_to_pair(be: Backend, t: Tensor) -> PairTensor: + """complex tensor -> PairTensor(re, im) of bf16.""" bf = _bf16_dtype() re = be.cast(be.real(t), bf) im = be.cast(be.imag(t), bf) - return be.stack([re, im], axis=-1) + return PairTensor(re, im) -def _pair_to_complex(be: Backend, pair: Tensor) -> Tensor: - """pair of bf16 -> complex tensor (recombine; no copy risk via cast).""" - re = be.cast(pair[..., 0], cons.rdtypestr) - im = be.cast(pair[..., 1], cons.rdtypestr) +def _pair_to_complex(be: Backend, pair: PairTensor) -> Tensor: + """PairTensor of bf16 -> complex tensor (recombine; no copy risk via cast).""" + re, im = pair.unpack() + re = be.cast(re, cons.rdtypestr) + im = be.cast(im, cons.rdtypestr) return be.cast(re + 1j * im, cons.dtypestr) def _pair_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: """Complex tensordot = 4 real bf16 tensordots (4M). Uses be.tensordot (never patched).""" - ar, ai = a[..., 0], a[..., 1] - br, bi = b[..., 0], b[..., 1] + ar, ai = PairTensor.unpack_pair(a) + br, bi = PairTensor.unpack_pair(b) cr = be.tensordot(ar, br, axes) - be.tensordot(ai, bi, axes) ci = be.tensordot(ar, bi, axes) + be.tensordot(ai, br, axes) - return be.stack([cr, ci], axis=-1) + return PairTensor.pack_result(be, cr, ci, not isinstance(a, PairTensor)) +# Strategy note: bf16 single-operand einsum manually decomposes into diagonal → +# sum → transpose (staying in bf16 end-to-end) because numpy's ``np.einsum`` +# rejects bfloat16 dtypes. This is different from the tropical algebra's +# ``_tropical_einsum`` single-operand path, which delegates to ``be.einsum`` +# for repeated-index resolution — the tropical backend (float64) has no such +# dtype restriction. Both produce equivalent results under their respective +# semirings, but the implementation strategy is dictated by dtype constraints +# rather than algebraic differences. def _einsum_single_operand_half( be: Backend, x: Tensor, lhs: str, out_subs: str ) -> Tensor: @@ -81,52 +93,23 @@ def _einsum_single_operand_half( return x -def _pair_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: - """Complex einsum = 4 real bf16 einsums (4M for 2 operands, 2 for 1). - - **Two-operand (genuine bf16):** manually decomposed into ``be.tensordot`` + - ``be.transpose``, because numpy's C ``einsum`` rejects ``ml_dtypes.bfloat16`` - (it is a standalone C routine with a hardcoded dtype allowlist, not a - ufunc). ``np.tensordot`` accepts bf16 because it dispatches through ufunc - loops, so the compute is genuine bf16 end-to-end — no float32 upcast. - The decomposition parses the einsum subscript equation to find contracted - axes, then uses ``tensordot`` for the contraction and ``transpose`` to - match the output subscript order. - - **Single-operand (genuine bf16):** decomposed into ``np.diagonal`` + - ``be.sum`` + ``be.transpose``, all bf16-safe. Handles reductions, - transposes, diagonals, and traces — the full einsum single-operand - semantics without any float32 upcast. - - ``_pair_tensordot`` needs no such routing because ``np.tensordot`` already - preserves bf16 (verified: it accumulates in bf16, not float32). - cotengra feeds only 1-2-operand equations here (it decomposes hyperedges - itself), all of which are pairwise contractions that map cleanly to - tensordot. - """ +def _pair_einsum(be: Backend, eq: str, *operands: Tensor) -> PairTensor: + """Complex einsum = 4 real bf16 einsums (4M for 2 operands, 2 for 1).""" if len(operands) == 1: - # ── Single-operand: pure bf16 (decompose into sum + transpose + diagonal) ── a = operands[0] - - # Implicit mode (no ``->``) is an identity / no-op at the einsum level. if "->" not in eq: return a - lhs, out_subs = eq.split("->") - return be.stack( - [ - _einsum_single_operand_half(be, a[..., 0], lhs, out_subs), - _einsum_single_operand_half(be, a[..., 1], lhs, out_subs), - ], - axis=-1, + ar, ai = PairTensor.unpack_pair(a) + return PairTensor( + _einsum_single_operand_half(be, ar, lhs, out_subs), + _einsum_single_operand_half(be, ai, lhs, out_subs), ) - # ── 2-operand: bf16-safe tensordot decomposition ────────────────── a, b = operands - ar, ai = a[..., 0], a[..., 1] - br, bi = b[..., 0], b[..., 1] + ar, ai = PairTensor.unpack_pair(a) + br, bi = PairTensor.unpack_pair(b) - # Parse the einsum equation once if "->" not in eq: raise ValueError( f"implicit-mode einsum {eq!r} not supported for bf16; use explicit '->'" @@ -143,14 +126,12 @@ def _pair_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: out_order = list(out_subs) def _contract(x: Tensor, y: Tensor) -> Tensor: - """Pairwise bf16-safe einsum → tensordot + optional transpose.""" if contracted: a_axes = [a_subs.index(c) for c in contracted] b_axes = [b_subs.index(c) for c in contracted] result = be.tensordot(x, y, axes=(a_axes, b_axes)) else: result = be.tensordot(x, y, axes=0) - free_order = a_free + b_free if free_order != out_order: perm = [free_order.index(c) for c in out_order] @@ -159,7 +140,7 @@ def _contract(x: Tensor, y: Tensor) -> Tensor: cr = _contract(ar, br) - _contract(ai, bi) ci = _contract(ar, bi) + _contract(ai, br) - return be.stack([cr, ci], axis=-1) + return PairTensor.pack_result(be, cr, ci, not isinstance(a, PairTensor)) class PairBf16Representation(Representation): @@ -188,9 +169,5 @@ def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: @contextlib.contextmanager def bcomplex32() -> Iterator[None]: - prev = cons.get_contraction_algebra() - cons.set_contraction_algebra(ComplexPairAlgebra()) - try: + with cons.runtime_contraction_algebra(ComplexPairAlgebra()): yield - finally: - cons.set_contraction_algebra(prev) diff --git a/applications/tropical_algebra.py b/applications/tropical_algebra.py index d6690e0d..cd7efb4d 100644 --- a/applications/tropical_algebra.py +++ b/applications/tropical_algebra.py @@ -34,7 +34,11 @@ import numpy as np import tensorcircuit.cons as cons -from tensorcircuit.contraction_algebra import ContractionAlgebra, Representation +from tensorcircuit.contraction_algebra import ( + ContractionAlgebra, + PairTensor, + Representation, +) Tensor = Any Backend = Any @@ -99,6 +103,11 @@ def _expand_to_layout( return be.reshape(tt, newshape) +# Strategy note: single-operand tropical einsum uses ``be.einsum`` for +# repeated-index resolution (diagonal gather). This is safe because tropical +# tensors are float64, which numpy's einsum accepts natively. For bf16 +# (which numpy rejects), see ``_einsum_single_operand_half`` in +# ``applications/bcomplex32_algebra.py`` for the manual decomposition approach. def _tropical_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: """max-plus einsum: product -> +, sum -> max. Handles 1- or 2-operand forms.""" if len(operands) == 1: @@ -163,37 +172,24 @@ def split_energy_count(stacked: Tensor) -> "tuple[Tensor, Tensor]": return arr[..., 0], arr[..., 1] -def _stack_last(be: Backend, x: Tensor, n: Tensor) -> Tensor: - """Stack two same-shape tensors along a new trailing axis (portable). - - All tc-ng concrete backends (numpy, jax, torch, tensorflow, cupy) implement - ``be.stack`` via ``tensorcircuit.backends.abstract_backend``, so the portable - path is used directly. ``be.stack`` takes a Python sequence and inserts a new - axis; ``axis=-1`` puts it last to match the ``[..., 2]`` convention. - """ - return be.stack([x, n], axis=-1) - - -def _counting_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: +def _counting_tensordot(be: Backend, a: Tensor, b: Tensor, axes: Any) -> PairTensor: """(max-plus energy, degeneracy count) pairwise contraction. - ``a`` and ``b`` are stacked ``[..., 2]`` tensors (``[..., 0]`` = energy, - ``[..., 1]`` = count). The energy stream follows max-plus; the count stream - sums ``a_count * b_count`` over the contracted positions that achieve the - energy max (within ``_EPS``), so each max-tie contributes its multiplicity. + ``a`` and ``b`` are ``PairTensor`` (energy, count) or legacy stacked tensors. """ - a3x, b3x, out_shape = _pair_layout(be, a[..., 0], b[..., 0], axes) + ae, an = PairTensor.unpack_pair(a) + b_e, b_n = PairTensor.unpack_pair(b) + a3x, b3x, out_shape = _pair_layout(be, ae, b_e, axes) s = a3x + b3x # (m, k, n) energy pair sums y = be.max(s, axis=1) # (m, n) max energy per output slot - # abs(s - y) < _EPS, broadcast across the contracted axis k: reshape y to (m, 1, n) m_n = tuple(int(d) for d in be.shape_tuple(y)) y_b = be.reshape(y, (m_n[0], 1, m_n[1])) mask = be.abs(s - y_b) < _EPS - a3n, b3n, _ = _pair_layout(be, a[..., 1], b[..., 1], axes) + a3n, b3n, _ = _pair_layout(be, an, b_n, axes) pn = a3n * b3n # (m, k, n) pairwise count products cy = be.reshape(be.sum(pn * mask, axis=1), out_shape) y_shaped = be.reshape(y, out_shape) - return _stack_last(be, y_shaped, cy) + return PairTensor.pack_result(be, y_shaped, cy, not isinstance(a, PairTensor)) def _insert_axis_of( @@ -217,25 +213,21 @@ def _resolve_repeats( """Resolve intra-operand repeated indices via per-stream diagonal gather. Pure per-stream indexing -- the energy and count streams are gathered - independently (``pair[..., 0]`` and ``pair[..., 1]``) so the trailing - ``[..., 2]`` stack axis is never treated as an index by ``be.einsum``. The - max/tie degeneracy logic lives entirely in the later reduction; this helper - only rewrites each stream's index layout. + independently via ``.unpack()``; the trailing pair axis stays invisible to + ``be.einsum``. The max/tie degeneracy logic lives entirely in the later + reduction; this helper only rewrites each stream's index layout. Returns ``(energy, count, resolved_idxs)`` where ``resolved_idxs`` has no - repeats. If ``idxs`` has no repeats, the streams are sliced out unchanged - and ``idxs`` is returned as-is (passthrough -- 2-operand behavior is - identical to the pre-helper code path). + repeats. If ``idxs`` has no repeats, the streams are returned unchanged + and ``idxs`` is returned as-is. """ if len(set(idxs)) != len(idxs): resolved = "".join(dict.fromkeys(idxs)) - # Diagonal extraction: exactly one element contributes per output position, - # so the standard (sum, multiply) einsum is equivalent to max — both - # reduce to identity for a singleton set of values. - e = be.einsum("".join(idxs) + "->" + resolved, pair[..., 0]) - n = be.einsum("".join(idxs) + "->" + resolved, pair[..., 1]) + e, n = PairTensor.unpack_pair(pair) + e = be.einsum("".join(idxs) + "->" + resolved, e) + n = be.einsum("".join(idxs) + "->" + resolved, n) return e, n, list(resolved) - return pair[..., 0], pair[..., 1], list(idxs) + return PairTensor.unpack_pair(pair) + (list(idxs),) def _counting_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: @@ -271,7 +263,9 @@ def _counting_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: remaining = [c for c in idxs if c not in contract] out_e = be.transpose(e, tuple(remaining.index(c) for c in rhs)) out_n = be.transpose(n, tuple(remaining.index(c) for c in rhs)) - return _stack_last(be, out_e, out_n) + return PairTensor.pack_result( + be, out_e, out_n, not isinstance(operands[0], PairTensor) + ) a, b = operands lhs, rhs = eq.split("->") ia_s, ib_s = lhs.split(",") @@ -298,18 +292,13 @@ def _counting_einsum(be: Backend, eq: str, *operands: Tensor) -> Tensor: remaining = [c for c in all_idx if c not in contract] out_x = be.transpose(sx, tuple(remaining.index(c) for c in out_idx)) out_n = be.transpose(pn, tuple(remaining.index(c) for c in out_idx)) - return _stack_last(be, out_x, out_n) + return PairTensor.pack_result(be, out_x, out_n, not isinstance(a, PairTensor)) class CountingRepresentation(Representation): """Attach count=1 (the counting-semiring multiplicative identity) to each - leaf; decode splits the final stacked ``[..., 2]`` tensor into energy - (primary) + count (aux). - - encode: ``t -> stack([t, ones_like(t, float64)], axis=-1)``. Per-tensor, - topology-agnostic, run once on the leaves before any pairwise contraction. - decode: ``tensor[..., 0]`` is the energy (primary, rank == len(output_set)); - ``tensor[..., 1]`` is the degeneracy count (aux, co-indexed with energy). + leaf via ``PairTensor(t, ones)``. Decode unpacks into energy (primary) + + count (aux, stashed in ``_last_aux`` for ``degeneracy()``). """ name = "counting" @@ -318,11 +307,13 @@ def encode(self, be: Backend, tensors: List[Tensor]) -> List[Tensor]: out = [] for t in tensors: ones = be.ones_like(t, dtype="float64") - out.append(be.stack([t, ones], axis=-1)) + out.append(PairTensor(t, ones)) return out def decode(self, be: Backend, tensor: Tensor) -> Tuple[Tensor, Dict[str, Tensor]]: - return tensor[..., 0], {"count": tensor[..., 1]} + e, n = PairTensor.unpack_pair(tensor) + self._last_aux = {"count": n} + return e, {} class CountingTropicalAlgebra(ContractionAlgebra): @@ -330,10 +321,7 @@ class CountingTropicalAlgebra(ContractionAlgebra): Carries ``CountingRepresentation``: encode attaches count=1 to each leaf; decode splits the final pair into energy (primary) + count (aux, stashed - via ``cons._stash_aux_outputs`` for ``degeneracy()`` to read). No - ``on_contraction_start`` hook is needed: the unconditional aux clear in - ``cons._algebraic_base_contraction`` already wipes stale state per - contraction (standard or non-standard). + in ``CountingRepresentation._last_aux`` for ``degeneracy()`` to read). """ name = "counting_maxplus" @@ -348,23 +336,15 @@ def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: def degeneracy() -> Optional[Any]: """Degeneracy (number of optimal configs) of the most recent counting - contraction. - - Call inside the ``counting_tropical()`` block (like - ``recover_configuration``) after a contraction has run. Returns the count - tensor (same shape as the primary energy) stashed by - ``CountingRepresentation.decode``; returns ``None`` if no counting decode - has stashed a count for the current contraction. - - The aux side-channel is cleared at the start of every contraction that - routes through ``cons._algebraic_base_contraction`` (the unconditional - ``_stash_aux_outputs({})`` at the top of that function), so a counting - contraction followed by another algebraic contraction wipes stale state. - Note that ``cons._base`` only routes to ``_algebraic_base_contraction`` - for non-standard algebras, hyperedge inputs, or ``use_primitives=True``; - a standard contraction on the original ``_base`` path does NOT clear aux. - """ - return cons._aux_outputs().get("count") + contraction. Call inside ``counting_tropical()`` after a contraction has run. + Returns ``None`` if no counting algebra is active.""" + alg = cons.get_contraction_algebra() + if alg is None: + return None + rep = alg.representation + if hasattr(rep, "_last_aux"): + return rep._last_aux.get("count") + return None # ===== Section 3: tracking + configuration recovery ===== @@ -768,31 +748,19 @@ def recover_configuration() -> Dict[Any, int]: def tropical(track: bool = False) -> Iterator[None]: """Contract under the max-plus (tropical) algebra within the block. - Swaps ``cons._contraction_algebra`` directly (no monkey-patch activation): - the in-source ``cons._base`` routes to ``_algebraic_base_contraction`` for - any non-standard algebra, so this is sufficient. - ``track=True`` switches to ``MaxPlusTrackingAlgebra`` so ``recover_configuration()`` can recover the optimal configuration after the contraction. Off by default -> zero behaviour change relative to plain max-plus. """ algebra = MaxPlusTrackingAlgebra() if track else MaxPlusAlgebra() - prev = cons.get_contraction_algebra() - cons.set_contraction_algebra(algebra) - try: + with cons.runtime_contraction_algebra(algebra): yield - finally: - cons.set_contraction_algebra(prev) @contextlib.contextmanager def counting_tropical() -> Iterator[None]: """Contract under the counting (energy, degeneracy) max-plus algebra within the block.""" - prev = cons.get_contraction_algebra() - cons.set_contraction_algebra(CountingTropicalAlgebra()) - try: + with cons.runtime_contraction_algebra(CountingTropicalAlgebra()): yield - finally: - cons.set_contraction_algebra(prev) diff --git a/docs/superpowers/plans/2026-07-19-contractor-kwargs.md b/docs/superpowers/plans/2026-07-19-contractor-kwargs.md deleted file mode 100644 index 5d010b01..00000000 --- a/docs/superpowers/plans/2026-07-19-contractor-kwargs.md +++ /dev/null @@ -1,173 +0,0 @@ -# Replace `prefer_einsum` with `get_contractor_kwargs()` — Implementation Plan - -> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. - -**Goal:** Replace the `prefer_einsum: bool` class attribute on `ContractionAlgebra` with a `get_contractor_kwargs()` method, keeping the same behavior while decoupling the ABC from cotengra internals. - -**Architecture:** A single new method on the ABC returns a dict of kwargs to forward to `ctg.core.make_contractor`. The default returns `{}`. `ComplexPairAlgebra` overrides to return `{"prefer_einsum": True}`. `cons.py` calls the method instead of reading the attribute. - -**Tech Stack:** Pure Python — no new dependencies. - -## Global Constraints - -- No behavior change for any algebra -- All existing tests must pass -- Three files touched: `base.py`, `bcomplex32_algebra.py`, `cons.py` - ---- - -### Task 1: Add `get_contractor_kwargs()` to ABC, remove `prefer_einsum` attribute - -**Files:** -- Modify: `tensorcircuit/contraction_algebra/base.py:59-69` - -**Interfaces:** -- Produces: `ContractionAlgebra.get_contractor_kwargs() -> dict` (default returns `{}`) - -- [ ] **Step 1: Replace `prefer_einsum` attribute with `get_contractor_kwargs()` method** - -Replace lines 59-69: - -```python - name: str = "abstract" - representation: Representation = IdentityRepresentation() - # When operands carry a trailing non-physical storage axis (e.g. the - # complex pair), cotengra's tensordot mode post-transposes results via - # autoray and mishandles that extra axis (ValueError: axes don't match array) - # -- set True to force einsum-only execution, which forwards operands verbatim - # to ``einsum`` and skips the autoray transpose. Default False keeps tensordot - # mode: tropical config-recovery backtracking depends on the tensordot - # intermediate layout and would break under forced einsum. - prefer_einsum: bool = False -``` - -With: - -```python - name: str = "abstract" - representation: Representation = IdentityRepresentation() - - def get_contractor_kwargs(self) -> dict: - """Extra kwargs forwarded to cotengra's ``make_contractor``. - - Override to return ``{'prefer_einsum': True}`` when your algebra's - ``tensordot`` kernel carries non-physical storage axes (e.g. the - complex pair axis) that cotengra's post-tensordot autoray - transpose would mishandle (ValueError: axes don't match array). - ``prefer_einsum=True`` forces einsum-only execution, which skips - the transpose entirely. - - Default ``{}`` keeps the standard tensordot+einsum mix — required - by tropical config-recovery backtracking, which depends on the - tensordot intermediate layout. - """ - return {} -``` - -- [ ] **Step 2: Verify the file is valid Python** - -Run: `D:/Software/miniconda3/envs/tcng/python.exe -c "from tensorcircuit.contraction_algebra.base import ContractionAlgebra; print(ContractionAlgebra().get_contractor_kwargs())"` -Expected: `{}` - -- [ ] **Step 3: Commit** - -```bash -git add tensorcircuit/contraction_algebra/base.py -git commit -m "refactor: replace prefer_einsum attr with get_contractor_kwargs() method on ContractionAlgebra" -``` - ---- - -### Task 2: Override `get_contractor_kwargs()` in `ComplexPairAlgebra` - -**Files:** -- Modify: `applications/bcomplex32_algebra.py:175-178` - -**Interfaces:** -- Consumes: `ContractionAlgebra.get_contractor_kwargs() -> dict` (from Task 1) - -- [ ] **Step 1: Replace `prefer_einsum` class attribute with method override** - -Replace line 178: - -```python - prefer_einsum = True # pair operands carry a trailing storage axis -``` - -With: - -```python - def get_contractor_kwargs(self) -> dict: - return {"prefer_einsum": True} -``` - -The class now looks like: - -```python -class ComplexPairAlgebra(ContractionAlgebra): - name = "bcomplex32_pair" - representation = PairBf16Representation() - - def get_contractor_kwargs(self) -> dict: - return {"prefer_einsum": True} - - def tensordot(self, be: Backend, a: Tensor, b: Tensor, axes: Any) -> Tensor: - ... -``` - -- [ ] **Step 2: Verify** - -Run: `D:/Software/miniconda3/envs/tcng/python.exe -c "from applications.bcomplex32_algebra import ComplexPairAlgebra; print(ComplexPairAlgebra().get_contractor_kwargs())"` -Expected: `{'prefer_einsum': True}` - -- [ ] **Step 3: Commit** - -```bash -git add applications/bcomplex32_algebra.py -git commit -m "refactor: use get_contractor_kwargs() override in ComplexPairAlgebra" -``` - ---- - -### Task 3: Update `cons.py` to call `get_contractor_kwargs()` - -**Files:** -- Modify: `tensorcircuit/cons.py:784-785` - -**Interfaces:** -- Consumes: `ContractionAlgebra.get_contractor_kwargs() -> dict` (from Task 1) - -- [ ] **Step 1: Replace `alg.prefer_einsum` with `**alg.get_contractor_kwargs()`** - -Replace lines 784-785: - -```python - contractor = ctg.core.make_contractor( - tree, implementation=impl, prefer_einsum=alg.prefer_einsum - ) -``` - -With: - -```python - contractor = ctg.core.make_contractor( - tree, implementation=impl, **alg.get_contractor_kwargs() - ) -``` - -- [ ] **Step 2: Verify no remaining references to `prefer_einsum` on algebra** - -Run: `cd "e:\Study\.AShare\OneDrive\OneDriveSync\session\tc\tensorcircuit-ng" && grep -rn "prefer_einsum" tensorcircuit/ applications/ --include="*.py" | grep -v test | grep -v __pycache__` -Expected: only `bcomplex32_algebra.py` (inside the new method returning the dict), and possibly `base.py` docstring mentioning it as an example - -- [ ] **Step 3: Run existing tests** - -Run: `D:/Software/miniconda3/envs/tcng/python.exe -m pytest tests/test_bcomplex32_algebra.py tests/test_contraction_algebra.py tests/test_tropical.py -v` -Expected: all pass - -- [ ] **Step 4: Commit** - -```bash -git add tensorcircuit/cons.py -git commit -m "refactor: consume alg.get_contractor_kwargs() in _algebraic_base_contraction" -``` \ No newline at end of file diff --git a/docs/superpowers/specs/2026-07-19-contractor-kwargs-design.md b/docs/superpowers/specs/2026-07-19-contractor-kwargs-design.md deleted file mode 100644 index c01955d0..00000000 --- a/docs/superpowers/specs/2026-07-19-contractor-kwargs-design.md +++ /dev/null @@ -1,80 +0,0 @@ -# ContractionAlgebra: replace `prefer_einsum` class attribute with `get_contractor_kwargs()` method - -**Date:** 2026-07-19 -**Status:** draft -**PR:** feat/contraction-algebra-tropical - -## Motivation - -`ContractionAlgebra.prefer_einsum` is a cotengra-specific escape hatch that leaks -implementation details into the algebra ABC. The reviewer asked: "but you can -customize the tensordot, can you avoid this?" — the flag exists **not** because -`alg.tensordot` itself is broken, but because cotengra's `Contractor` applies a -post-`tensordot` autoray transpose whose `perm` doesn't know about non-physical -storage axes (e.g., the bf16 pair axis). The custom `tensordot` cannot prevent -this transpose, so `prefer_einsum=True` is the only way to skip it. - -However, `prefer_einsum` should not be a first-class ABC attribute. Instead, let -the algebra declare **contractor-level options** through a method whose return -value is transparently forwarded. - -## Design - -### ABC: `get_contractor_kwargs() -> dict` - -```python -class ContractionAlgebra(ABC): - def get_contractor_kwargs(self) -> dict: - """Extra kwargs forwarded to cotengra's make_contractor. - - Override to return ``{'prefer_einsum': True}`` when your kernel - carries non-physical storage axes that cotengra's post-tensordot - autoray transpose cannot handle. - """ - return {} -``` - -- Remove `prefer_einsum: bool = False`. -- Keep the explanatory comment, moved into the method docstring. - -### `ComplexPairAlgebra`: override - -```python -class ComplexPairAlgebra(ContractionAlgebra): - def get_contractor_kwargs(self) -> dict: - return {"prefer_einsum": True} -``` - -Remove the `prefer_einsum = True` class attribute. - -### `cons.py`: consume the method - -```python -# Before: -contractor = ctg.core.make_contractor( - tree, implementation=impl, prefer_einsum=alg.prefer_einsum -) - -# After: -contractor = ctg.core.make_contractor( - tree, implementation=impl, **alg.get_contractor_kwargs() -) -``` - -## Files changed - -| File | Change | -|------|--------| -| `tensorcircuit/contraction_algebra/base.py` | Remove `prefer_einsum` attr; add `get_contractor_kwargs()` method | -| `applications/bcomplex32_algebra.py` | Replace `prefer_einsum = True` with `get_contractor_kwargs()` override | -| `tensorcircuit/cons.py` | Call `alg.get_contractor_kwargs()` instead of `alg.prefer_einsum` | - -## Non-goals - -- Does not change cotengra's internal behavior -- Does not add new escape hatches — only re-packages the existing one -- No behavior change for `StandardAlgebra` or `TropicalAlgebra` - -## Reviewer reply - -> Good catch — `prefer_einsum` is now hidden behind a `get_contractor_kwargs()` method on the ABC so the algebra itself doesn't expose cotengra-specific flags as class attributes. \ No newline at end of file diff --git a/docs/superpowers/specs/2026-07-19-prefer-einsum-contractor-kwargs-design.md b/docs/superpowers/specs/2026-07-19-prefer-einsum-contractor-kwargs-design.md deleted file mode 100644 index 972f59be..00000000 --- a/docs/superpowers/specs/2026-07-19-prefer-einsum-contractor-kwargs-design.md +++ /dev/null @@ -1,83 +0,0 @@ -# Extract `prefer_einsum` from ContractionAlgebra ABC into `get_contractor_kwargs()` - -Date: 2026-07-19 - -## Motivation - -`ContractionAlgebra` currently carries a `prefer_einsum: bool` class attribute whose -sole purpose is to be forwarded to `ctg.core.make_contractor(tree, prefer_einsum=...)`. -It exists because cotengra's tensordot mode applies a post-contraction transpose whose -permutation axes do not account for non-physical storage axes (e.g. the trailing -complex pair axis), causing `ValueError: axes don't match array`. - -This flag is an implementation detail of *how* cotengra executes contractions and -does not belong in the algebra ABC. The reviewer correctly observed that since the -algebra already customizes `tensordot`, a second escape-hatch flag seems redundant -— but the actual problem is in cotengra's internal transpose, not in our tensordot. - -## Design - -Replace the hard attribute with a method: - -```python -class ContractionAlgebra(ABC): - def get_contractor_kwargs(self) -> Dict[str, Any]: - """Extra keyword arguments forwarded to ``ctg.core.make_contractor``. - - Override to return ``{"prefer_einsum": True}`` when your contraction - kernels produce tensors with non-physical storage axes (e.g. the - complex pair axis) that cotengra's post-tensordot transpose - cannot handle. - - The default empty dict preserves standard tensordot mode, which - tropical config-recovery backtracking depends on. - """ - return {} -``` - -`ComplexPairAlgebra` overrides it: - -```python -class ComplexPairAlgebra(ContractionAlgebra): - def get_contractor_kwargs(self) -> Dict[str, Any]: - return {"prefer_einsum": True} -``` - -In `cons.py` the dispatch changes from: - -```python -contractor = ctg.core.make_contractor( - tree, implementation=impl, prefer_einsum=alg.prefer_einsum -) -``` - -to: - -```python -contractor = ctg.core.make_contractor( - tree, implementation=impl, **alg.get_contractor_kwargs() -) -``` - -## Trade-offs - -- **Pro:** `ContractionAlgebra` ABC no longer exposes cotengra's `prefer_einsum` - parameter as a first-class attribute. -- **Pro:** The dict return type is extensible — future cotengra kwargs - (e.g. `strip_exponent`) can be supplied without changing the ABC again. -- **Pro:** Default `{}` keeps standard algebras unchanged; tropical contracts - through the same path as before. -- **Con:** The method signature still implies that the algebra knows about - cotengra's constructor interface. This coupling is acceptable because the - algebra's *only* consumer is `_algebraic_base_contraction`, which builds a - cotengra contractor. - -## Affected files - -| File | Change | -|------|--------| -| `tensorcircuit/contraction_algebra/base.py` | Replace `prefer_einsum: bool` attribute with `get_contractor_kwargs()` method | -| `applications/bcomplex32_algebra.py` | Replace `prefer_einsum = True` with `get_contractor_kwargs()` override | -| `tensorcircuit/cons.py` | Call `alg.get_contractor_kwargs()` instead of reading `alg.prefer_einsum` | - -Tests do not need changes — the runtime behavior is identical. \ No newline at end of file diff --git a/tensorcircuit/cons.py b/tensorcircuit/cons.py index 6a2ac07d..bc12507e 100644 --- a/tensorcircuit/cons.py +++ b/tensorcircuit/cons.py @@ -157,32 +157,6 @@ def set_contraction_algebra(alg: Optional[_ContractionAlgebra]) -> None: _contraction_algebra = alg -# Aux output side-channel for multi-output (non-standard) algebras. Cleared and -# refilled on each non-standard contraction by ``_stash_aux_outputs``; the -# primary tensor is returned through the normal contraction return value, while -# aux (e.g. counting degeneracy, sharing the primary's physical axis order) is -# stashed here for the caller to read via ``_aux_outputs``. -# -# Thread-safety: this global, along with ``_trace`` and ``_ctx`` in -# ``applications/tropical_algebra.py``, is module-level mutable state cleared -# per contraction. Concurrent contractions will corrupt each other. Acceptable -# for the single-threaded research use case; see also -# ``docs/superpowers/plans/2026-07-19-L3-bf16-backend-generic.md``. -_aux_outputs_store: Dict[str, Any] = {} - - -def _stash_aux_outputs(aux: Dict[str, Any]) -> None: - """Stash the most recent contraction's aux outputs (side channel for - multi-output algebras, e.g. counting degeneracy). Cleared per contraction.""" - _aux_outputs_store.clear() - _aux_outputs_store.update(aux) - - -def _aux_outputs() -> Dict[str, Any]: - """Test/accessor: snapshot of the stashed aux outputs.""" - return dict(_aux_outputs_store) - - def set_function_backend(backend: Optional[str] = None) -> Callable[..., Any]: """ Function decorator to set function-level runtime backend @@ -747,15 +721,6 @@ def _wrap_omeco_optimizer(optimizer: Any) -> Any: return optimizer -def _stash_permuted_aux( - aux: Any, output_edge_order: Any, dangling_edges: Any, kbe: Any -) -> None: - """Apply the output_edge_order permutation to aux (count co-indexed with energy).""" - order = output_edge_order if output_edge_order is not None else dangling_edges - perm = [dangling_edges.index(e) for e in order] - _stash_aux_outputs({k: kbe.transpose(v, tuple(perm)) for k, v in aux.items()}) - - def _run_contraction( raw_tensors: Any, input_sets: Any, @@ -763,18 +728,32 @@ def _run_contraction( size_dict: Any, algorithm: Any, ns: bool, - alg: Any, - kbe: Any, + algebra: Any, be: Any, strip_exponent: bool, ctg: Any, ) -> Any: - """Run the (possibly single-tensor) contraction, returning (final, exponent).""" + """Run the (possibly single-tensor) contraction, returning ``(final, exponent)``. + + Parameters: + raw_tensors: leaf tensors extracted from the tn.Node topology. + input_sets: cotengra input index sets (one per leaf). + output_set: cotengra output index set. + size_dict: cotengra size dictionary (index → dimension). + algorithm: cotengra path-finder (e.g. ``opt_einsum.paths.dynamic_programming``). + ns: ``True`` when a non-standard ``ContractionAlgebra`` is active. + algebra: the active ``ContractionAlgebra`` instance (carries kernels + hooks). + be: backend for all tensor operations (algebra kernels + tn.Node wrap). + strip_exponent: if ``True``, return the exponent separately via cotengra. + ctg: the ``cotengra`` module (imported once at the call site). + """ exponent = 0.0 if len(raw_tensors) == 1: # Avoid cotengra bug for empty contraction paths eq = input_sets[0] + "->" + output_set - final = alg.einsum(kbe, eq, *raw_tensors) if ns else be.einsum(eq, *raw_tensors) + final = ( + algebra.einsum(be, eq, *raw_tensors) if ns else be.einsum(eq, *raw_tensors) + ) else: path = algorithm(input_sets, output_set, size_dict) logger.info("the contraction path is given as %s" % str(path)) @@ -782,13 +761,13 @@ def _run_contraction( input_sets, output_set, size_dict, path=path ) if ns: - alg.on_contractor_ready(tree) + algebra.on_contractor_ready(tree) impl = ( - functools.partial(alg.einsum, kbe), - functools.partial(alg.tensordot, kbe), + functools.partial(algebra.einsum, be), + functools.partial(algebra.tensordot, be), ) contractor = ctg.core.make_contractor( - tree, implementation=impl, **alg.get_contractor_kwargs() + tree, implementation=impl, **algebra.get_contractor_kwargs() ) final = contractor(*raw_tensors) elif not strip_exponent: @@ -800,17 +779,20 @@ def _run_contraction( return final, exponent -def _decode_aux(ns: bool, rep: Any, kbe: Any, final: Any, output_set: Any) -> Any: - """Decode the contraction output under a non-standard algebra; else no aux.""" +def _decode(ns: bool, rep: Any, be: Any, final: Any, output_set: Any) -> Any: + """Decode the contraction output under a non-standard algebra. + + ``rep.decode`` is responsible for stashing any aux internally.""" if not ns: - return final, {} - primary, aux = rep.decode(kbe, final) - assert primary.ndim == len(output_set), ( - "representation.decode primary rank %d != len(output_set) %d; " - "decode must strip non-physical storage axes before tn.Node wraps it" - % (primary.ndim, len(output_set)) - ) - return primary, aux + return final + primary, _ = rep.decode(be, final) + if primary.ndim != len(output_set): + raise ValueError( + "representation.decode primary rank %d != len(output_set) %d; " + "decode must strip non-physical storage axes before tn.Node wraps it" + % (primary.ndim, len(output_set)) + ) + return primary def _rewire_dangling_edges(final_node: Any, dangling_edges: Any, nodes: Any) -> None: @@ -838,26 +820,20 @@ def _algebraic_base_contraction( import cotengra as ctg raw_tensors, input_sets, output_set, size_dict = _extract_topology(nodes) - be = nodes[0].backend # tn backend: standard native ops + tn.Node wrap alg = get_contraction_algebra() if alg is not None: rep = alg.representation - kbe = backend # algebra kernels need the tc backend (max/argmax/...) else: - rep = IdentityRepresentation() # no-op; _decode_aux skips it when ns=False - kbe = be - # be and backend are normally the same object: tn.set_default_backend syncs them. - - _stash_aux_outputs({}) + rep = IdentityRepresentation() # no-op; _decode skips it when ns=False if alg is not None: if kws.get("strip_exponent", False): raise ValueError( "strip_exponent is incompatible with a non-standard ContractionAlgebra" ) alg.on_contraction_start(nodes) - raw_tensors = rep.encode(kbe, raw_tensors) + raw_tensors = rep.encode(backend, raw_tensors) final, exponent = _run_contraction( raw_tensors, @@ -866,16 +842,15 @@ def _algebraic_base_contraction( size_dict, algorithm, alg is not None, - alg, - kbe, - be, - kws.get("strip_exponent", False), - ctg, + algebra=alg, + be=backend, + strip_exponent=kws.get("strip_exponent", False), + ctg=ctg, ) - final, aux = _decode_aux(alg is not None, rep, kbe, final, output_set) + final = _decode(alg is not None, rep, backend, final, output_set) - final_node = tn.Node(final, backend=be) + final_node = tn.Node(final, backend=backend) # Resolve dangling edges in the same order as in _extract_topology dangling_edges = sorted_edges(tn.get_subgraph_dangling(nodes)) @@ -889,8 +864,12 @@ def _algebraic_base_contraction( final_node.reorder_edges(list(output_edge_order)) # Apply the same output_edge_order permutation to aux (count co-indexed with energy) - if alg is not None and aux: - _stash_permuted_aux(aux, output_edge_order, dangling_edges, kbe) + if alg is not None and hasattr(rep, "_last_aux") and rep._last_aux: + order = output_edge_order if output_edge_order is not None else dangling_edges + perm = [dangling_edges.index(e) for e in order] + rep._last_aux = { + k: backend.transpose(v, tuple(perm)) for k, v in rep._last_aux.items() + } if kws.get("strip_exponent", False) and alg is None: return final_node, exponent @@ -1054,8 +1033,6 @@ def _base( # ORIGINAL EXECUTION PATH (100% Backward Compatible) # ========================================== - _stash_aux_outputs({}) # clear stale aux from prior non-standard contraction - for edge in edges: if not edge.is_disabled: # if its disabled we already contracted it if edge.is_trace(): @@ -1438,13 +1415,11 @@ def wrapper(f: Callable[..., Any]) -> Callable[..., Any]: @wraps(f) def newf(*args: Any, **kws: Any) -> Any: old_contractor = getattr(thismodule, "contractor") - old_algebra = get_contraction_algebra() set_contractor(*confargs, **confkws) try: return f(*args, **kws) finally: _set_global_contractor(old_contractor) - set_contraction_algebra(old_algebra) return newf @@ -1460,13 +1435,26 @@ def runtime_contractor(*confargs: Any, **confkws: Any) -> Iterator[Any]: :rtype: Iterator[Any] """ old_contractor = getattr(thismodule, "contractor") - old_algebra = get_contraction_algebra() nc = set_contractor(*confargs, **confkws) try: yield nc finally: _set_global_contractor(old_contractor) - set_contraction_algebra(old_algebra) + + +@contextmanager +def runtime_contraction_algebra(alg: _ContractionAlgebra) -> Iterator[None]: + """Context manager to temporarily set a non-standard contraction algebra. + + Mirrors ``runtime_backend`` / ``runtime_dtype``: saves the current algebra, + sets ``alg`` for the block, restores on exit. + """ + prev = get_contraction_algebra() + set_contraction_algebra(alg) + try: + yield + finally: + set_contraction_algebra(prev) def split_rules( diff --git a/tensorcircuit/contraction_algebra.py b/tensorcircuit/contraction_algebra.py index 4e7cfc79..82ed2a2f 100644 --- a/tensorcircuit/contraction_algebra.py +++ b/tensorcircuit/contraction_algebra.py @@ -12,6 +12,8 @@ (bf16 pair). """ +from __future__ import annotations + from abc import ABC, abstractmethod from typing import Any, Dict, List, Tuple @@ -103,9 +105,52 @@ def einsum(self, be: Backend, eq: str, *operands: Tensor) -> Tensor: return be.einsum(eq, *operands) +class PairTensor: + """Virtual tensor wrapping two halves (e.g., re/im, energy/count). + + Algebra kernels unpack via ``.unpack()`` and operate on the halves + directly, then return new ``PairTensor`` results. ``encode()`` and + ``decode()`` are the only places that create / consume the pair + representation. The pair axis is kept off ``tn.Node`` (cotengra sees only + the primary half's topology via the symbolic tree), so no ``.shape`` / + ``.ndim`` array-protocol is exposed. + """ + + __slots__ = ("_p", "_s") + + def __init__(self, primary: Tensor, secondary: Tensor): + self._p = primary + self._s = secondary + + def unpack(self) -> Tuple[Tensor, Tensor]: + return self._p, self._s + + @staticmethod + def unpack_pair(pair: "PairTensor | Tensor") -> Tuple[Tensor, Tensor]: + """Unpack a PairTensor or legacy stack [..., 2] — backward compatible.""" + if isinstance(pair, PairTensor): + return pair.unpack() + return pair[..., 0], pair[..., 1] + + @staticmethod + def pack_result( + be: Any, primary: Tensor, secondary: Tensor, legacy: bool + ) -> "PairTensor | Tensor": + """Return PairTensor or legacy stack matching input format.""" + # Kernels intentionally accept both PairTensor and legacy stacked + # [..., 2] operands and echo the input format on output, so direct + # callers and the counting tests that build be.stack([..., -1]) keep + # working. The contraction flow itself is PairTensor end-to-end (encode + # wraps every leaf); only direct kernel entry points see the legacy form. + if legacy: + return be.stack([primary, secondary], axis=-1) + return PairTensor(primary, secondary) + + __all__ = [ "ContractionAlgebra", "StandardAlgebra", "Representation", "IdentityRepresentation", + "PairTensor", ] diff --git a/tests/test_bcomplex32_algebra.py b/tests/test_bcomplex32_algebra.py index ad0e56d6..d262dc27 100644 --- a/tests/test_bcomplex32_algebra.py +++ b/tests/test_bcomplex32_algebra.py @@ -99,10 +99,11 @@ def test_pair_einsum_keeps_bfloat16_dtype(): bf = ml_dtypes.bfloat16 a = np.array([[1.0 + 2.0j, 3.0j], [-1.0j, 2.0 - 1.0j]], dtype=np.complex64) b = np.array([[0.5 + 0.5j, 1.0j], [2.0j, -1.0 + 1.0j]], dtype=np.complex64) - pa = _complex_to_pair(be, a) # bf16 pair, shape (2, 2, 2) + pa = _complex_to_pair(be, a) pb = _complex_to_pair(be, b) - out = np.asarray(_pair_einsum(be, "ij,jk->ik", pa, pb)) - assert out.dtype == bf, f"_pair_einsum upcast to {out.dtype}; expected bfloat16" + result = _pair_einsum(be, "ij,jk->ik", pa, pb) + re, _ = result.unpack() + assert re.dtype == bf, f"_pair_einsum upcast to {re.dtype}; expected bfloat16" def test_bf16_ghz8_runs_and_matches_native(): diff --git a/tests/test_contraction_algebra.py b/tests/test_contraction_algebra.py index efe64467..ec4a19b1 100644 --- a/tests/test_contraction_algebra.py +++ b/tests/test_contraction_algebra.py @@ -175,7 +175,7 @@ def einsum(self, be, eq, *ops): prev = cons.get_contraction_algebra() cons.set_contraction_algebra(BadAlg()) try: - with pytest.raises(AssertionError): + with pytest.raises(ValueError): cons._algebraic_base_contraction( [a, b], algorithm=opt_einsum.paths.dynamic_programming, @@ -185,11 +185,15 @@ def einsum(self, be, eq, *ops): cons.set_contraction_algebra(prev) -def test_aux_outputs_stashed_and_reordered(): - # The full aux-reorder test is Task 8 (counting). Here we only confirm the - # side-channel store API exists and round-trips a value. - cons._stash_aux_outputs({"count": 5}) - assert cons._aux_outputs()["count"] == 5 +def test_counting_representation_stashes_aux(): + import numpy as np + from applications.tropical_algebra import CountingRepresentation + + rep = CountingRepresentation() + t = np.array([[3.0, 7.0], [1.0, 4.0]]) + primary, _ = rep.decode(None, t) + assert np.array_equal(primary, np.array([3.0, 1.0])) + assert np.array_equal(rep._last_aux["count"], np.array([7.0, 4.0])) def test_representation_identity_roundtrip(): @@ -369,11 +373,15 @@ def test_aux_stash_handles_ignore_edge_order_with_none_order(monkeypatch): import opt_einsum class AuxRep(Representation): + name = "aux" + _last_aux = {} + def encode(self, be, tensors): return tensors def decode(self, be, tensor): - return tensor, {"count": np.ones_like(tensor)} # non-empty aux + self._last_aux = {"count": np.ones_like(tensor)} + return tensor, {} class AuxAlg(ContractionAlgebra): name = "aux" @@ -399,7 +407,7 @@ def einsum(self, be, eq, *ops): output_edge_order=None, ignore_edge_order=True, ) - assert "count" in cons._aux_outputs() # aux stashed without crashing + assert "count" in cons.get_contraction_algebra().representation._last_aux finally: cons.set_contraction_algebra(prev) @@ -408,3 +416,55 @@ def test_get_contractor_kwargs_default(): assert StandardAlgebra().get_contractor_kwargs() == {} # Concrete algebra without override inherits the default assert _NS().get_contractor_kwargs() == {} + + +# --- Coverage: single-tensor path + the strip_exponent x algebra guard. +# Both hit lines in cons._algebraic_base_contraction / _run_contraction that no +# other test exercises (len(raw_tensors)==1 cotengra empty-path workaround, and +# the strip_exponent incompatibility guard). + + +def test_single_tensor_contraction_uses_algebra_einsum(): + # A one-node contraction takes the len(raw_tensors)==1 branch, bypassing + # cotengra; it must still dispatch through algebra.einsum under a + # non-standard algebra (here StandardAlgebra), not just be.einsum. + import tensornetwork as tn + import opt_einsum + + t = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + a = tn.Node(t) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(StandardAlgebra()) + try: + node = cons._algebraic_base_contraction( + [a], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[a[0], a[1]], + ) + # "ab->ab" identity contraction: the result equals the leaf tensor. + np.testing.assert_allclose(np.array(node.tensor), t) + finally: + cons.set_contraction_algebra(prev) + + +def test_strip_exponent_incompatible_with_algebra_raises(): + # strip_exponent=True under any non-standard algebra is rejected up front. + import tensornetwork as tn + import opt_einsum + import pytest + + a = tn.Node(np.array([1.0, 2.0])) + b = tn.Node(np.array([3.0, 4.0])) + tn.connect(a[0], b[0]) + prev = cons.get_contraction_algebra() + cons.set_contraction_algebra(StandardAlgebra()) + try: + with pytest.raises(ValueError, match="strip_exponent is incompatible"): + cons._algebraic_base_contraction( + [a, b], + algorithm=opt_einsum.paths.dynamic_programming, + output_edge_order=[], + strip_exponent=True, + ) + finally: + cons.set_contraction_algebra(prev)