clean up cons

Author: refraction-ray

.agents/memory/workflow_and_quality.md

@@ -40,6 +40,7 @@ Use this file for changes that touch public API, tests, docs, or code quality.
 - Benchmark external libraries in isolated environments and match precision/settings before making performance claims.
 - Library plotting helpers should accept an optional `ax` and should not call `plt.show()` internally.
 - In read-only or sandboxed environments, redirect caches such as `NUMBA_CACHE_DIR`, `MPLCONFIGDIR`, and mypy's cache directory in the shell command or local harness used for validation, not in tracked source files, unless the repo explicitly needs that behavior.
+- If a test or import failure plausibly comes from sandbox restrictions or third-party runtime cache/process-pool setup, rerun the exact command with escalated permissions before attributing the failure to TensorCircuit logic; sandbox artifacts can mask the real failure mode.
 - If `cotengra` fails at import time with `ImportError: cannot import name 'get_namespace' from autoray`, treat it as an environment dependency mismatch and upgrade `autoray` before debugging TensorCircuit's contraction code.
 - On macOS sandboxed runs, `conda run -n <env> ...` can silently resolve to the base interpreter; verify `sys.executable` or call the env's Python directly when reproducing environment-sensitive issues.
 - Cotengra hyper-optimization that reaches `joblib`/`loky` process-pool setup can fail inside the sandbox with `PermissionError` from `os.sysconf("SC_SEM_NSEMS_MAX")`; for end-to-end benchmark validation, rerun the exact benchmark command with escalated permissions rather than changing TensorCircuit logic.

examples/cotengra_visualize_path.py

@@ -25,7 +25,7 @@ def main():
 
     print(f"Created a {n}-qubit circuit with {len(c._nodes)} gates.")
 
-    # 2. Extract algebraic contraction topology
+    # 2. Extract a stable contraction topology description for cotengra
     # We use the internal `_nodes` which represents the state vector contraction
     info, _ = tc.cons.get_tn_info(c._nodes)
     inputs, output, size_dict = info

examples/ng_whitepaper/README.md

@@ -0,0 +1,7 @@
+# TensorCircuit-NG Whitepaper Numerical Demonstrations
+
+This directory contains the core numerical demonstration scripts, benchmarks, and figures accompanying the TensorCircuit-NG whitepaper:
+
+**Paper Link**: [TensorCircuit-NG: Next-Generation Tensor Network and Differentiable Quantum Simulation](https://arxiv.org/abs/2602.14167) (arXiv:2602.14167)
+
+Each script is named according to the whitepaper section it belongs to and is fully executable, demonstrating the power, flexibility, and performance of the framework.

examples/reproduce_papers/README.md

@@ -0,0 +1,19 @@
+# Reproduce Papers with TensorCircuit-NG
+
+This directory contains high-fidelity, reproducible simulations of key quantum computing and quantum information papers using **TensorCircuit-NG**. Each subfolder represents a specific paper reproduction.
+
+## Directory Structure
+
+For any new paper reproduction, follow the standardized naming and folder structure:
+```
+examples/reproduce_papers/
+└── <YYYY>_<keywords>/            # e.g., 2026_diff_qec_surface/
+    ├── meta.yaml                 # Standard metadata describing the reproduction strategy
+    ├── main.py                   # Main executable reproduction script (JAX-native, JIT-friendly)
+    └── outputs/                  # Directory containing generated results and plots
+        └── result.png            # Reproduced figure/plot
+```
+
+## Guidance for AI Agents
+
+Refer to the complete instruction set in the `arxiv-reproduce` skill (`.agents/skills/arxiv-reproduce/SKILL.md`) for detailed specifications.

tensorcircuit/cons.py

@@ -80,6 +80,12 @@ def sorted_edges(edges: Iterator[tn.Edge]) -> List[tn.Edge]:
 # these above lines are just for mypy, it is not very good at evaluating runtime object
 
 
+def _set_global_contractor(contractor_fn: Callable[..., Any]) -> None:
+    for module in sys.modules:
+        if module.startswith(package_name):
+            setattr(sys.modules[module], "contractor", contractor_fn)
+
+
 def set_tensornetwork_backend(
     backend: Optional[str] = None, set_global: bool = True
 ) -> Any:
@@ -1008,9 +1014,7 @@ def runtime_nodes_capture(key: str = "nodes") -> Iterator[Any]:
     except NodesReturn as e:
         captured_value[key] = e.value
     finally:
-        for module in sys.modules:
-            if module.startswith(package_name):
-                setattr(sys.modules[module], "contractor", old_contractor)
+        _set_global_contractor(old_contractor)
 
 
 def custom(
@@ -1023,6 +1027,8 @@ def custom(
     use_primitives: Optional[bool] = None,
     **kws: Any,
 ) -> Any:
+    local_kws = dict(kws)
+    debug_level = local_kws.pop("debug_level", debug_level)
     if len(nodes) < 5:
         alg = opt_einsum.paths.optimal
         # not good at minimize WRITE actually...
@@ -1033,19 +1039,18 @@ def custom(
             ignore_edge_order,
             debug_level=debug_level,
             use_primitives=use_primitives,
-            **kws,
+            **local_kws,
         )
 
     total_size = None
     has_hyperedges = any(isinstance(n, tn.CopyNode) for n in nodes)
-    if kws.get("preprocessing", None) and not has_hyperedges:
+    if local_kws.get("preprocessing", None) and not has_hyperedges:
         # nodes = _full_light_cone_cancel(nodes)
         nodes, total_size = _merge_single_gates(nodes)
     if not isinstance(optimizer, list):
         alg = partial(optimizer, memory_limit=memory_limit)
     else:
         alg = optimizer
-    debug_level = kws.get("debug_level", 0)
     return _base(
         nodes,
         alg,
@@ -1054,7 +1059,7 @@ def custom(
         total_size,
         debug_level=debug_level,
         use_primitives=use_primitives,
-        **kws,
+        **local_kws,
     )
 
 
@@ -1068,42 +1073,27 @@ def custom_stateful(
     use_primitives: Optional[bool] = None,
     **kws: Any,
 ) -> Any:
-    if len(nodes) < 5:
-        alg = opt_einsum.paths.optimal
-        # dynamic_programming has a potential bug for outer product
-        # not good at minimize WRITE actually...
-        return _base(
-            nodes,
-            alg,
-            output_edge_order,
-            ignore_edge_order,
-            use_primitives=use_primitives,
-        )
-
-    total_size = None
-    has_hyperedges = any(isinstance(n, tn.CopyNode) for n in nodes)
-    if kws.get("preprocessing", None) and not has_hyperedges:
-        nodes, total_size = _merge_single_gates(nodes)
     if opt_conf is None:
         opt_conf = {}
     opt = optimizer(**opt_conf)  # reinitiate the optimizer each time
     if kws.get("contraction_info", None):
         opt = contraction_info_decorator(opt)
-    alg = partial(opt, memory_limit=memory_limit)
-    debug_level = kws.get("debug_level", 0)
-
-    return _base(
+    local_kws = dict(kws)
+    debug_level = local_kws.pop("debug_level", 0)
+    local_kws.pop("contraction_info", None)
+    return custom(
         nodes,
-        alg,
-        output_edge_order,
-        ignore_edge_order,
-        total_size,
+        opt,
+        memory_limit=memory_limit,
+        output_edge_order=output_edge_order,
+        ignore_edge_order=ignore_edge_order,
         debug_level=debug_level,
         use_primitives=use_primitives,
+        **local_kws,
     )
 
 
-# only work for custom
+# used by custom contractor variants
 def contraction_info_decorator(algorithm: Callable[..., Any]) -> Callable[..., Any]:
     """Decorator to add contraction information logging to an optimizer.
 
@@ -1249,9 +1239,7 @@ def set_contractor(
             **kws,
         )
     if set_global:
-        for module in sys.modules:
-            if module.startswith(package_name):
-                setattr(sys.modules[module], "contractor", cf)
+        _set_global_contractor(cf)
     return cf
 
 
@@ -1273,11 +1261,10 @@ def wrapper(f: Callable[..., Any]) -> Callable[..., Any]:
         def newf(*args: Any, **kws: Any) -> Any:
             old_contractor = getattr(thismodule, "contractor")
             set_contractor(*confargs, **confkws)
-            r = f(*args, **kws)
-            for module in sys.modules:
-                if module.startswith(package_name):
-                    setattr(sys.modules[module], "contractor", old_contractor)
-            return r
+            try:
+                return f(*args, **kws)
+            finally:
+                _set_global_contractor(old_contractor)
 
         return newf
 
@@ -1294,10 +1281,10 @@ def runtime_contractor(*confargs: Any, **confkws: Any) -> Iterator[Any]:
     """
     old_contractor = getattr(thismodule, "contractor")
     nc = set_contractor(*confargs, **confkws)
-    yield nc
-    for module in sys.modules:
-        if module.startswith(package_name):
-            setattr(sys.modules[module], "contractor", old_contractor)
+    try:
+        yield nc
+    finally:
+        _set_global_contractor(old_contractor)
 
 
 def split_rules(

tests/test_circuit.py

@@ -2126,3 +2126,24 @@ def test_strip_exponent_no_hyperedge(backend, reset_contractor):
     res_node, exponent = result
     np.testing.assert_allclose(tc.backend.numpy(res_node.tensor), 1.0, atol=1e-5)
     np.testing.assert_allclose(exponent, 10.0, atol=1e-5)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("jaxb")])
+def test_custom_stateful_strip_exponent(backend, reset_contractor):
+    tc.set_contractor(
+        "custom_stateful",
+        optimizer=oem.RandomGreedy,
+        max_repeats=1,
+        strip_exponent=True,
+    )
+
+    nodes = [
+        tn.Node(tc.backend.convert_to_tensor(10.0, tc.rdtypestr)) for _ in range(6)
+    ]
+
+    result = tc.cons.contractor(nodes)
+    assert isinstance(result, tuple)
+
+    res_node, exponent = result
+    np.testing.assert_allclose(tc.backend.numpy(res_node.tensor), 1.0, atol=1e-5)
+    np.testing.assert_allclose(exponent, 6.0, atol=1e-5)

tests/test_quantum.py

@@ -745,7 +745,8 @@ def test_qop2quimb(backend):
     ket_inds_mps = [f"k{i}" for i in range(nwires_mps)]
     vec_from_quimb = np.ravel(quimb_mps.to_dense(ket_inds_mps))
 
-    np.testing.assert_allclose(vec_from_qop, vec_from_quimb, atol=1e-5)
+    atol = 5e-5 if tc.backend.name == "jax" else 1e-5
+    np.testing.assert_allclose(vec_from_qop, vec_from_quimb, atol=atol)
 
 
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])