赛前稳定版

tools/baseline_mps_expectation.py

@@ -19,6 +19,22 @@ from qibotn.backends.qmatchatea import QMatchaTeaBackend
 from qibotn.backends.vidal_tebd import run_vidal_ring_xz
 
 
+def optional_int(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return int(text)
+
+
+def optional_float(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return float(text)
+
+
+def format_optional(value, fmt="g"):
+    return "None" if value is None else format(value, fmt)
+
+
 def build_circuit(nqubits, nlayers, seed):
     return build_benchmark_circuit("brickwall_cnot", nqubits, nlayers, seed)
 
@@ -35,7 +51,8 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--nqubits", type=int, default=40)
     parser.add_argument("--nlayers", type=int, default=30)
-    parser.add_argument("--bond", "--bonds", dest="bond", type=int, default=512)
+    parser.add_argument("--bond", "--bonds", dest="bond", type=optional_int, default=512)
+    parser.add_argument("--cut-ratio", type=optional_float, default=1e-12)
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--tensor-module", choices=("numpy", "torch"), default="torch")
     parser.add_argument("--torch-threads", type=int, default=32)
@@ -109,7 +126,8 @@ def main():
             mpi_label = f"MPIMPS/{size}" if args.mpi_ct else "SR"
         print(
             f"nqubits={args.nqubits} nlayers={args.nlayers} "
-            f"bond={args.bond} seed={args.seed} "
+            f"bond={format_optional(args.bond)} "
+            f"cut_ratio={format_optional(args.cut_ratio)} seed={args.seed} "
             f"tensor_module={args.tensor_module} svd_control=E! "
             f"compile_circuit=True mpi={mpi_label} executor={args.executor}"
         )
@@ -128,7 +146,7 @@ def main():
             value, timings = run_segment_vidal_mpi_ring_xz(
                 circuit,
                 max_bond=args.bond,
-                cut_ratio=1e-12,
+                cut_ratio=args.cut_ratio,
                 tensor_module=args.tensor_module,
                 comm=MPI.COMM_WORLD,
             )
@@ -136,7 +154,7 @@ def main():
             value = run_vidal_ring_xz(
                 circuit,
                 max_bond=args.bond,
-                cut_ratio=1e-12,
+                cut_ratio=args.cut_ratio,
                 tensor_module=args.tensor_module,
             )
     else:
@@ -144,7 +162,7 @@ def main():
         backend.configure_tn_simulation(
             ansatz="MPS",
             max_bond_dimension=args.bond,
-            cut_ratio=1e-12,
+            cut_ratio=args.cut_ratio,
             svd_control="E!",
             tensor_module=args.tensor_module,
             compile_circuit=True,

tools/example_tn_case.py

@@ -0,0 +1,33 @@
+"""Example custom case for tools/run_tn_custom.py."""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+from qibo import Circuit, gates
+
+
+def build_circuit(nqubits, nlayers, seed):
+    rng = np.random.default_rng(seed)
+    circuit = Circuit(nqubits)
+    for layer in range(nlayers):
+        for qubit in range(nqubits):
+            circuit.add(gates.RY(qubit, theta=rng.uniform(-math.pi, math.pi)))
+            circuit.add(gates.RZ(qubit, theta=rng.uniform(-math.pi, math.pi)))
+        for qubit in range(layer % 2, nqubits - 1, 2):
+            circuit.add(gates.RXX(qubit, qubit + 1, theta=rng.uniform(-0.7, 0.7)))
+            circuit.add(gates.RZZ(qubit, qubit + 1, theta=rng.uniform(-0.7, 0.7)))
+    return circuit
+
+
+def build_observable(nqubits, seed):
+    return {
+        "terms": [
+            {
+                "coefficient": 1.0 / max(1, nqubits - 1),
+                "operators": [("Z", site), ("Z", site + 1)],
+            }
+            for site in range(nqubits - 1)
+        ]
+    }

tools/inspect_contraction_tree.py

@@ -0,0 +1,208 @@
+"""Inspect cotengra contraction trees for dominant torch matmul shapes."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import math
+import pickle
+from collections import Counter, defaultdict
+from pathlib import Path
+
+
+def _prod(values):
+    out = 1
+    for value in values:
+        out *= int(value)
+    return out
+
+
+def _broadcast_batch(a_batch, b_batch):
+    if a_batch == b_batch:
+        return _prod(a_batch)
+    if not a_batch:
+        return _prod(b_batch)
+    if not b_batch:
+        return _prod(a_batch)
+
+    ndim = max(len(a_batch), len(b_batch))
+    a_batch = (1,) * (ndim - len(a_batch)) + tuple(a_batch)
+    b_batch = (1,) * (ndim - len(b_batch)) + tuple(b_batch)
+    return _prod(max(a, b) for a, b in zip(a_batch, b_batch))
+
+
+def _load_tree(path, index):
+    with Path(path).open("rb") as f:
+        payload = pickle.load(f)
+    trees = payload["trees"] if isinstance(payload, dict) else payload
+    if not isinstance(trees, (list, tuple)):
+        trees = [trees]
+    return trees[index]
+
+
+def _analyze_tree(tree):
+    contract_mod = importlib.import_module("cotengra.contract")
+    contractions = contract_mod.extract_contractions(tree)
+    size_dict = tree.size_dict
+    ops = []
+    counts = Counter()
+
+    for op_index, (parent, left, right, tdot, arg, perm) in enumerate(contractions):
+        if left is None and right is None:
+            counts["preprocess"] += 1
+            continue
+
+        left_inds = tree.get_inds(left)
+        right_inds = tree.get_inds(right)
+        parent_inds = tree.get_inds(parent)
+        left_shape = tuple(size_dict[ix] for ix in left_inds)
+        right_shape = tuple(size_dict[ix] for ix in right_inds)
+
+        if tdot:
+            parsed = contract_mod._parse_tensordot_axes_to_matmul(
+                arg,
+                left_shape,
+                right_shape,
+            )
+        else:
+            parsed = contract_mod._parse_eq_to_batch_matmul(
+                arg,
+                left_shape,
+                right_shape,
+            )
+
+        (
+            _eq_a,
+            _eq_b,
+            new_shape_a,
+            new_shape_b,
+            _new_shape_ab,
+            _perm_ab,
+            pure_multiplication,
+        ) = parsed
+
+        matmul_shape = None
+        matmul_flops = 0
+        if pure_multiplication:
+            kind = "mul"
+        else:
+            a_shape = tuple(new_shape_a or left_shape)
+            b_shape = tuple(new_shape_b or right_shape)
+            batch = _broadcast_batch(a_shape[:-2], b_shape[:-2])
+            m, k, n = int(a_shape[-2]), int(a_shape[-1]), int(b_shape[-1])
+            kind = "mm" if batch == 1 else "bmm"
+            matmul_shape = (batch, m, k, n)
+            matmul_flops = batch * m * k * n
+
+        tree_flops = int(tree.get_flops(parent))
+        out_size = int(tree.get_size(parent))
+        ops.append(
+            {
+                "index": op_index,
+                "kind": kind,
+                "matmul_shape": matmul_shape,
+                "matmul_flops": matmul_flops,
+                "tree_flops": tree_flops,
+                "out_size": out_size,
+                "left_shape": left_shape,
+                "right_shape": right_shape,
+                "left_rank": len(left_inds),
+                "right_rank": len(right_inds),
+                "out_rank": len(parent_inds),
+                "perm": perm,
+            }
+        )
+        counts[kind] += 1
+
+    return contractions, ops, counts
+
+
+def _format_log(value, base):
+    return "-inf" if value <= 0 else f"{math.log(value, base):.3f}"
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("tree", help="Pickle file containing one tree or {'trees': [...]}.")
+    parser.add_argument("--index", type=int, default=0, help="Tree index in the file.")
+    parser.add_argument("--top", type=int, default=20, help="Number of top ops to print.")
+    parser.add_argument(
+        "--dtype-bytes",
+        type=int,
+        default=8,
+        help="Bytes per element for memory estimates, for example 8 for complex64.",
+    )
+    args = parser.parse_args()
+
+    tree = _load_tree(args.tree, args.index)
+    contractions, ops, counts = _analyze_tree(tree)
+    nslices = int(getattr(tree, "multiplicity", 1))
+    per_slice_flops = sum(op["tree_flops"] for op in ops)
+    per_slice_write = sum(op["out_size"] for op in ops)
+    max_out = max((op["out_size"] for op in ops), default=0)
+    all_flops = per_slice_flops * nslices
+    all_write = per_slice_write * nslices
+
+    print(f"tree={args.tree} index={args.index}")
+    print(
+        "summary "
+        f"slices={nslices} contractions={len(contractions)} "
+        f"counts={dict(counts)}"
+    )
+    print(
+        "per_slice "
+        f"log10_flops={_format_log(per_slice_flops, 10)} "
+        f"log10_write={_format_log(per_slice_write, 10)} "
+        f"log2_max_output={_format_log(max_out, 2)} "
+        f"max_output_gib={max_out * args.dtype_bytes / 1024**3:.6g}"
+    )
+    print(
+        "all_slices "
+        f"log10_flops={_format_log(all_flops, 10)} "
+        f"log10_write={_format_log(all_write, 10)}"
+    )
+
+    print(f"\ntop_{args.top}_ops_by_flops")
+    for op in sorted(ops, key=lambda item: item["tree_flops"], reverse=True)[: args.top]:
+        print(
+            f"op={op['index']} kind={op['kind']} "
+            f"flops={op['tree_flops']:.6e} out={op['out_size']:.6e} "
+            f"matmul={op['matmul_shape']} "
+            f"ranks=({op['left_rank']},{op['right_rank']}->{op['out_rank']}) "
+            f"lhs={op['left_shape']} rhs={op['right_shape']}"
+        )
+
+    by_shape = defaultdict(lambda: [0, 0, 0])
+    for op in ops:
+        shape = op["matmul_shape"]
+        if shape is None:
+            continue
+        by_shape[shape][0] += 1
+        by_shape[shape][1] += op["tree_flops"]
+        by_shape[shape][2] += op["out_size"]
+
+    print(f"\ntop_{args.top}_matmul_shapes_by_flops")
+    for shape, (count, flops, out_size) in sorted(
+        by_shape.items(),
+        key=lambda item: item[1][1],
+        reverse=True,
+    )[: args.top]:
+        print(
+            f"shape={shape} count={count} "
+            f"flops={flops:.6e} output={out_size:.6e}"
+        )
+
+    print(f"\ntop_{args.top}_matmul_shapes_by_count")
+    for shape, (count, flops, out_size) in sorted(
+        by_shape.items(),
+        key=lambda item: item[1][0],
+        reverse=True,
+    )[: args.top]:
+        print(
+            f"shape={shape} count={count} "
+            f"flops={flops:.6e} output={out_size:.6e}"
+        )
+
+
+if __name__ == "__main__":
+    main()

tools/manage_tn_dask_cluster.sh

@@ -0,0 +1,223 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Manage the dask cluster used by TN path search.
+#
+# Defaults target two servers:
+#   scheduler: 10.20.1.103:8786
+#   workers:   10.20.1.103, 10.20.1.102
+#
+# Usage:
+#   tools/manage_tn_dask_cluster.sh start
+#   tools/manage_tn_dask_cluster.sh status
+#   tools/manage_tn_dask_cluster.sh stop
+#
+# Common overrides:
+#   SCHEDULER_HOST=10.20.1.103
+#   WORKER_HOSTS="10.20.1.103 10.20.1.102"
+#   NWORKERS=48
+#   NTHREADS=1
+#   ROOT_DIR=/home/yx/qibotn
+#   PYTHON_BIN=.venv/bin/python
+
+ROOT_DIR="${ROOT_DIR:-/home/yx/qibotn}"
+PYTHON_BIN="${PYTHON_BIN:-.venv/bin/python}"
+SCHEDULER_HOST="${SCHEDULER_HOST:-10.20.1.103}"
+SCHEDULER_PORT="${SCHEDULER_PORT:-8786}"
+DASHBOARD_ADDRESS="${DASHBOARD_ADDRESS:-:8787}"
+WORKER_HOSTS="${WORKER_HOSTS:-10.20.1.103 10.20.1.102}"
+NWORKERS="${NWORKERS:-48}"
+NTHREADS="${NTHREADS:-1}"
+MEMORY_LIMIT="${MEMORY_LIMIT:-0}"
+LOCAL_DIRECTORY="${LOCAL_DIRECTORY:-/tmp/qibotn-dask}"
+LOG_DIR="${LOG_DIR:-$ROOT_DIR/logs/dask}"
+SSH_BIN="${SSH_BIN:-ssh}"
+DASK_WORKER_TTL="${DASK_WORKER_TTL:-24 hours}"
+DASK_TICK_LIMIT="${DASK_TICK_LIMIT:-30 minutes}"
+DASK_LOST_WORKER_TIMEOUT="${DASK_LOST_WORKER_TIMEOUT:-30 minutes}"
+
+SCHEDULER_ADDR="tcp://${SCHEDULER_HOST}:${SCHEDULER_PORT}"
+
+is_local_host() {
+  local host="$1"
+  [[ "$host" == "localhost" || "$host" == "127.0.0.1" ]] && return 0
+  [[ "$host" == "$(hostname)" ]] && return 0
+  [[ "$host" == "$(hostname -f 2>/dev/null || true)" ]] && return 0
+  hostname -I 2>/dev/null | tr ' ' '\n' | grep -qx "$host"
+}
+
+run_on_host() {
+  local host="$1"
+  shift
+  local cmd="$*"
+  if is_local_host "$host"; then
+    bash -lc "$cmd"
+  else
+    "$SSH_BIN" "$host" "bash -lc $(printf '%q' "$cmd")"
+  fi
+}
+
+start_scheduler() {
+  local host="$SCHEDULER_HOST"
+  local log="$LOG_DIR/scheduler_${SCHEDULER_HOST}_${SCHEDULER_PORT}.log"
+  local pid_file="$LOG_DIR/scheduler_${SCHEDULER_HOST}_${SCHEDULER_PORT}.pid"
+  run_on_host "$host" "
+    set -euo pipefail
+    cd '$ROOT_DIR'
+    mkdir -p '$LOG_DIR'
+    if [[ -s '$pid_file' ]]; then
+      pid=\$(cat '$pid_file')
+      if kill -0 \"\$pid\" 2>/dev/null; then
+        echo \"scheduler already running on $host pid=\$pid\"
+        exit 0
+      fi
+    fi
+    DASK_DISTRIBUTED__SCHEDULER__WORKER_TTL='$DASK_WORKER_TTL' \
+    DASK_DISTRIBUTED__ADMIN__TICK__LIMIT='$DASK_TICK_LIMIT' \
+    DASK_DISTRIBUTED__DEPLOY__LOST_WORKER_TIMEOUT='$DASK_LOST_WORKER_TIMEOUT' \
+    setsid '$PYTHON_BIN' -m distributed.cli.dask_scheduler \
+      --host '$SCHEDULER_HOST' \
+      --port '$SCHEDULER_PORT' \
+      --dashboard-address '$DASHBOARD_ADDRESS' \
+      > '$log' 2>&1 < /dev/null &
+    pid=\$!
+    echo \"\$pid\" > '$pid_file'
+    echo \"scheduler host=$host pid=\$pid addr=$SCHEDULER_ADDR log=$log\"
+  "
+}
+
+start_worker() {
+  local host="$1"
+  local log="$LOG_DIR/worker_${host}.log"
+  local pid_file="$LOG_DIR/worker_${host}.pid"
+  run_on_host "$host" "
+    set -euo pipefail
+    cd '$ROOT_DIR'
+    mkdir -p '$LOG_DIR' '$LOCAL_DIRECTORY'
+    if [[ -s '$pid_file' ]]; then
+      pid=\$(cat '$pid_file')
+      if kill -0 \"\$pid\" 2>/dev/null; then
+        echo \"worker already running on $host pid=\$pid\"
+        exit 0
+      fi
+    fi
+    TCM_ENABLE=1 \
+    DASK_DISTRIBUTED__SCHEDULER__WORKER_TTL='$DASK_WORKER_TTL' \
+    DASK_DISTRIBUTED__ADMIN__TICK__LIMIT='$DASK_TICK_LIMIT' \
+    DASK_DISTRIBUTED__DEPLOY__LOST_WORKER_TIMEOUT='$DASK_LOST_WORKER_TIMEOUT' \
+    setsid '$PYTHON_BIN' -m distributed.cli.dask_worker \
+      '$SCHEDULER_ADDR' \
+      --host '$host' \
+      --nworkers '$NWORKERS' \
+      --nthreads '$NTHREADS' \
+      --memory-limit '$MEMORY_LIMIT' \
+      --local-directory '$LOCAL_DIRECTORY' \
+      > '$log' 2>&1 < /dev/null &
+    pid=\$!
+    echo \"\$pid\" > '$pid_file'
+    echo \"worker host=$host pid=\$pid scheduler=$SCHEDULER_ADDR log=$log\"
+  "
+}
+
+stop_host() {
+  local host="$1"
+  local scheduler_pid_file="$LOG_DIR/scheduler_${SCHEDULER_HOST}_${SCHEDULER_PORT}.pid"
+  local worker_pid_file="$LOG_DIR/worker_${host}.pid"
+  run_on_host "$host" "
+    set +e
+    for pid_file in '$worker_pid_file' '$scheduler_pid_file'; do
+      [[ -f \"\$pid_file\" ]] || continue
+      if [[ \"\$pid_file\" == '$scheduler_pid_file' && '$host' != '$SCHEDULER_HOST' ]]; then
+        continue
+      fi
+      pid=\$(cat \"\$pid_file\")
+      kill \"\$pid\" 2>/dev/null || true
+      rm -f \"\$pid_file\"
+    done
+    pkill -f '[d]istributed.cli.dask_worker.*$SCHEDULER_ADDR'
+    pkill -f '[d]istributed.cli.dask_scheduler.*--port $SCHEDULER_PORT'
+    true
+  "
+}
+
+status_host() {
+  local host="$1"
+  local scheduler_pid_file="$LOG_DIR/scheduler_${SCHEDULER_HOST}_${SCHEDULER_PORT}.pid"
+  local worker_pid_file="$LOG_DIR/worker_${host}.pid"
+  echo "--------------------------------------------------------------------------------"
+  echo "host=$host"
+  run_on_host "$host" "
+    set +e
+    for pid_file in '$worker_pid_file' '$scheduler_pid_file'; do
+      [[ -f \"\$pid_file\" ]] || continue
+      if [[ \"\$pid_file\" == '$scheduler_pid_file' && '$host' != '$SCHEDULER_HOST' ]]; then
+        continue
+      fi
+      pid=\$(cat \"\$pid_file\")
+      if kill -0 \"\$pid\" 2>/dev/null; then
+        ps -p \"\$pid\" -o pid,ppid,stat,etime,cmd --no-headers
+      else
+        echo \"stale pid_file=\$pid_file pid=\$pid\"
+      fi
+    done
+    pgrep -af '[d]istributed.cli.dask' || true
+  "
+}
+
+case "${1:-help}" in
+  start)
+    start_scheduler
+    sleep 2
+    for host in $WORKER_HOSTS; do
+      start_worker "$host"
+    done
+    echo
+    echo "Dask scheduler: $SCHEDULER_ADDR"
+    echo "Dashboard: http://$SCHEDULER_HOST$DASHBOARD_ADDRESS"
+    ;;
+  stop)
+    for host in $WORKER_HOSTS; do
+      stop_host "$host"
+    done
+    stop_host "$SCHEDULER_HOST"
+    ;;
+  status)
+    status_host "$SCHEDULER_HOST"
+    for host in $WORKER_HOSTS; do
+      [[ "$host" == "$SCHEDULER_HOST" ]] && continue
+      status_host "$host"
+    done
+    ;;
+  restart)
+    "$0" stop
+    sleep 2
+    "$0" start
+    ;;
+  help|*)
+    cat <<EOF
+Usage: tools/manage_tn_dask_cluster.sh [start|stop|restart|status]
+
+Defaults:
+  SCHEDULER_HOST=$SCHEDULER_HOST
+  SCHEDULER_PORT=$SCHEDULER_PORT
+  WORKER_HOSTS="$WORKER_HOSTS"
+  NWORKERS=$NWORKERS
+  NTHREADS=$NTHREADS
+  ROOT_DIR=$ROOT_DIR
+  PYTHON_BIN=$PYTHON_BIN
+  DASK_WORKER_TTL="$DASK_WORKER_TTL"
+  DASK_TICK_LIMIT=$DASK_TICK_LIMIT
+  DASK_LOST_WORKER_TIMEOUT=$DASK_LOST_WORKER_TIMEOUT
+
+Search command after start:
+  TCM_ENABLE=1 python -u tools/tn_contest_runner.py search \\
+    --case main1 \\
+    --dask-address $SCHEDULER_ADDR \\
+    --torch-threads 48 \\
+    --dtype complex64 \\
+    --tn-search-repeats 2048 \\
+    --tn-search-time 300
+EOF
+    exit 2
+    ;;
+esac

tools/mps_contest_runner.py

@@ -0,0 +1,313 @@
+#!/usr/bin/env python
+"""Contest-style multi-node Vidal/MPS expectation runner."""
+
+from __future__ import annotations
+
+import argparse
+import math
+import sys
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+from mpi4py import MPI
+from qibo import Circuit, gates, hamiltonians
+from qibo.symbols import X, Y, Z
+
+ROOT = Path(__file__).resolve().parents[1]
+SRC = ROOT / "src"
+if str(SRC) not in sys.path:
+    sys.path.insert(0, str(SRC))
+
+from qibotn.backends.vidal import VidalBackend  # noqa: E402
+from qibotn.expectation_runner import exact_for_observable  # noqa: E402
+
+
+@dataclass(frozen=True)
+class CaseSpec:
+    circuit_kind: str
+    observables: tuple[str, ...]
+    nqubits: int
+    nlayers: int
+    bond: int | None
+    seed: int
+
+
+CASES = {
+    "main1": CaseSpec(
+        circuit_kind="reversed_cnot",
+        observables=("ring_xz",),
+        nqubits=128,
+        nlayers=24,
+        bond=512,
+        seed=31001,
+    ),
+    "main2": CaseSpec(
+        circuit_kind="rxx_rzz",
+        observables=("open_zz", "range2_xx", "mixed_local"),
+        nqubits=128,
+        nlayers=32,
+        bond=1024,
+        seed=31002,
+    ),
+    "strong": CaseSpec(
+        circuit_kind="scramble",
+        observables=("ring_xz", "long_z_string", "dense3_spread"),
+        nqubits=256,
+        nlayers=48,
+        bond=2048,
+        seed=41001,
+    ),
+}
+
+
+def optional_int(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return int(text)
+
+
+def optional_float(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return float(text)
+
+
+def format_optional(value, fmt="g"):
+    return "None" if value is None else format(value, fmt)
+
+
+def set_torch_threads(nthreads):
+    try:
+        import torch
+
+        torch.set_num_threads(nthreads)
+    except Exception:
+        pass
+
+
+def add_single_qubit_layer(circuit, nqubits, rng, include_rx=False):
+    for qubit in range(nqubits):
+        circuit.add(gates.RY(qubit, theta=rng.uniform(-math.pi, math.pi)))
+        circuit.add(gates.RZ(qubit, theta=rng.uniform(-math.pi, math.pi)))
+        if include_rx:
+            circuit.add(gates.RX(qubit, theta=rng.uniform(-math.pi, math.pi)))
+
+
+def build_circuit(kind, nqubits, nlayers, seed):
+    rng = np.random.default_rng(seed)
+    circuit = Circuit(nqubits)
+
+    for layer in range(nlayers):
+        if kind == "reversed_cnot":
+            add_single_qubit_layer(circuit, nqubits, rng)
+            for qubit in range(0, nqubits - 1, 2):
+                gate = gates.CNOT(qubit + 1, qubit) if layer % 2 else gates.CNOT(qubit, qubit + 1)
+                circuit.add(gate)
+            for qubit in range(1, nqubits - 1, 2):
+                gate = gates.CNOT(qubit + 1, qubit) if layer % 2 == 0 else gates.CNOT(qubit, qubit + 1)
+                circuit.add(gate)
+
+        elif kind == "rxx_rzz":
+            add_single_qubit_layer(circuit, nqubits, rng, include_rx=True)
+            for qubit in range(layer % 2, nqubits - 1, 2):
+                circuit.add(gates.RXX(qubit, qubit + 1, theta=rng.uniform(-0.9, 0.9)))
+                circuit.add(gates.RZZ(qubit, qubit + 1, theta=rng.uniform(-0.9, 0.9)))
+
+        elif kind == "scramble":
+            add_single_qubit_layer(circuit, nqubits, rng, include_rx=True)
+            for qubit in range(layer % 2, nqubits - 1, 2):
+                circuit.add(gates.RXX(qubit, qubit + 1, theta=rng.uniform(-0.8, 0.8)))
+                circuit.add(gates.RZZ(qubit, qubit + 1, theta=rng.uniform(-0.8, 0.8)))
+                if layer % 5 == 4:
+                    circuit.add(gates.SWAP(qubit, qubit + 1))
+
+        else:
+            raise ValueError(f"Unknown circuit kind {kind!r}.")
+
+    return circuit
+
+
+def dense_observable(nqubits, qubits, seed, dim):
+    del nqubits
+    rng = np.random.default_rng(seed)
+    raw = rng.normal(size=(dim, dim)) + 1j * rng.normal(size=(dim, dim))
+    matrix = (raw + raw.conj().T) / 2.0
+    matrix = matrix / np.linalg.norm(matrix)
+    return {"matrix": matrix, "qubits": list(qubits)}
+
+
+def observable(kind, nqubits, seed):
+    q1 = nqubits // 4
+    q2 = nqubits // 2
+    q3 = (3 * nqubits) // 4
+    last = nqubits - 1
+
+    if kind == "boundary_ZZ_q1":
+        return hamiltonians.SymbolicHamiltonian(form=Z(q1 - 1) * Z(q1))
+    if kind == "boundary_ZZ_q2":
+        return hamiltonians.SymbolicHamiltonian(form=Z(q2 - 1) * Z(q2))
+    if kind == "boundary_ZZ_q3":
+        return hamiltonians.SymbolicHamiltonian(form=Z(q3 - 1) * Z(q3))
+    if kind == "long_Z_5_sites":
+        return hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(q1) * Z(q2) * Z(q3) * Z(last))
+    if kind == "mixed_XZYZX":
+        return hamiltonians.SymbolicHamiltonian(form=X(0) * Z(q1) * Y(q2) * Z(q3) * X(last))
+    if kind == "ring_xz":
+        form = 0
+        for qubit in range(nqubits):
+            form += 0.5 * X(qubit) * Z((qubit + 1) % nqubits)
+        return hamiltonians.SymbolicHamiltonian(form=form)
+    if kind == "open_zz":
+        form = 0
+        for qubit in range(nqubits - 1):
+            form += (1.0 / max(1, nqubits - 1)) * Z(qubit) * Z(qubit + 1)
+        return hamiltonians.SymbolicHamiltonian(form=form)
+    if kind == "range2_xx":
+        form = 0
+        for qubit in range(nqubits - 2):
+            form += (1.0 / max(1, nqubits - 2)) * X(qubit) * X(qubit + 2)
+        return hamiltonians.SymbolicHamiltonian(form=form)
+    if kind == "mixed_local":
+        form = 0.25 * X(0) - 0.5 * Z(last) + 0.125 * X(q1) * Z(q2) * Y(q3)
+        return hamiltonians.SymbolicHamiltonian(form=form)
+    if kind == "complex_iZ0":
+        return hamiltonians.SymbolicHamiltonian(form=1.0j * Z(0))
+    if kind == "dense2_mid":
+        return dense_observable(nqubits, (q2 - 1, q2), seed + 101, 4)
+    if kind == "dense3_spread":
+        return dense_observable(nqubits, (q1, q2, q3), seed + 202, 8)
+    raise ValueError(f"Unknown observable kind {kind!r}.")
+
+
+def selected_observables(args, case):
+    if args.observables:
+        return tuple(args.observables)
+    if args.obs_filter:
+        return tuple(x.strip() for x in args.obs_filter.split(",") if x.strip())
+    return case.observables
+
+
+def apply_case_defaults(args):
+    case = CASES[args.case]
+    if args.nqubits is None:
+        args.nqubits = case.nqubits
+    if args.nlayers is None:
+        args.nlayers = case.nlayers
+    if args.bond == "case-default":
+        args.bond = case.bond
+    if args.seed is None:
+        args.seed = case.seed
+    args.observables = selected_observables(args, case)
+
+
+def run_case(args):
+    set_torch_threads(args.torch_threads)
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    size = comm.Get_size()
+
+    case = CASES[args.case]
+    circuit = build_circuit(case.circuit_kind, args.nqubits, args.nlayers, args.seed)
+
+    if rank == 0:
+        print("=" * 88, flush=True)
+        print(
+            "backend=vidal_mps "
+            f"case={args.case} circuit={case.circuit_kind} ranks={size} "
+            f"nqubits={args.nqubits} nlayers={args.nlayers} gates={len(circuit.queue)} "
+            f"bond={format_optional(args.bond)} cut_ratio={format_optional(args.cut_ratio)} "
+            f"torch_threads={args.torch_threads} seed={args.seed} "
+            f"observables={','.join(args.observables)}",
+            flush=True,
+        )
+        print("observable exact value abs_error rel_error seconds trunc_sum trunc_max status", flush=True)
+
+    for obs_name in args.observables:
+        obs = observable(obs_name, args.nqubits, args.seed)
+        exact = None
+        if args.exact and rank == 0:
+            if args.nqubits > args.exact_max_qubits:
+                raise ValueError(
+                    f"--exact is limited to {args.exact_max_qubits} qubits by default."
+                )
+            exact = exact_for_observable(circuit, obs, args.nqubits)
+
+        backend = VidalBackend()
+        backend.configure_tn_simulation(
+            max_bond_dimension=args.bond,
+            cut_ratio=args.cut_ratio,
+            tensor_module="torch",
+            mpi_approach="CT",
+            mpi_num_procs=size,
+            fallback=False,
+        )
+
+        comm.Barrier()
+        start = time.perf_counter()
+        try:
+            value = backend.expectation(
+                circuit,
+                obs,
+                preprocess=True,
+                compile_circuit=False,
+            )
+            status = "ok"
+        except Exception as exc:
+            value = np.nan
+            status = type(exc).__name__ + ":" + str(exc).split("\n", 1)[0]
+        seconds = time.perf_counter() - start
+
+        if rank == 0:
+            abs_error = float("nan") if exact is None else abs(value - exact)
+            rel_error = float("nan") if exact is None else abs_error / max(abs(exact), 1e-15)
+            exact_text = "nan" if exact is None else f"{exact:.16e}"
+            print(
+                f"{obs_name} {exact_text} {value!r} "
+                f"{abs_error:.6e} {rel_error:.6e} {seconds:.3f} "
+                f"{backend.last_truncation_error:.6e} "
+                f"{backend.last_max_truncation_error:.6e} {status}",
+                flush=True,
+            )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("mode", choices=("run", "validate", "list"))
+    parser.add_argument("--case", choices=sorted(CASES), default="main1")
+    parser.add_argument("--observables", nargs="+")
+    parser.add_argument("--obs-filter", default="")
+    parser.add_argument("--nqubits", type=int)
+    parser.add_argument("--nlayers", type=int)
+    parser.add_argument("--bond", "--bonds", dest="bond", default="case-default")
+    parser.add_argument("--cut-ratio", type=optional_float, default=1e-12)
+    parser.add_argument("--seed", type=int)
+    parser.add_argument("--torch-threads", type=int, default=8)
+    parser.add_argument("--exact", action="store_true")
+    parser.add_argument("--exact-max-qubits", type=int, default=24)
+    args = parser.parse_args()
+
+    if args.mode == "list":
+        for name, case in CASES.items():
+            print(
+                f"{name}: circuit={case.circuit_kind} "
+                f"observables={','.join(case.observables)} "
+                f"nqubits={case.nqubits} nlayers={case.nlayers} "
+                f"bond={case.bond} seed={case.seed}"
+            )
+        return
+
+    apply_case_defaults(args)
+    if isinstance(args.bond, str):
+        args.bond = optional_int(args.bond)
+
+    if args.mode == "validate":
+        args.exact = True
+        args.nqubits = min(args.nqubits, args.exact_max_qubits)
+
+    run_case(args)
+
+
+if __name__ == "__main__":
+    main()

tools/run_tn_custom.py

@@ -0,0 +1,243 @@
+#!/usr/bin/env python
+"""Run TN expectation for a user-provided circuit and observable.
+
+The case module should define:
+
+    def build_circuit(nqubits, nlayers, seed): ...
+    def build_observable(nqubits, seed): ...
+
+``build_observable`` may return a Qibo SymbolicHamiltonian/form or the qibotn
+dict form:
+
+    {"terms": [
+        {"coefficient": 1.0, "operators": [("X", 0), ("Z", 1)]},
+    ]}
+
+For a single repeated Pauli string, pass ``--pauli-pattern`` instead of
+defining ``build_observable``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import inspect
+import json
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+SRC = ROOT / "src"
+if str(SRC) not in sys.path:
+    sys.path.insert(0, str(SRC))
+
+from qibotn.expectation_runner import (  # noqa: E402
+    ExpectationConfig,
+    exact_for_observable,
+    run_cpu_expectation,
+)
+
+
+def optional_int(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return int(text)
+
+
+def optional_float(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return float(text)
+
+
+def load_module(path):
+    path = Path(path).resolve()
+    spec = importlib.util.spec_from_file_location(path.stem, path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Cannot import case module from {path}.")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def call_builder(fn, **kwargs):
+    sig = inspect.signature(fn)
+    if any(p.kind == p.VAR_KEYWORD for p in sig.parameters.values()):
+        return fn(**kwargs)
+    accepted = {
+        name: value
+        for name, value in kwargs.items()
+        if name in sig.parameters
+    }
+    return fn(**accepted)
+
+
+def load_observable(args, module):
+    if args.pauli_pattern:
+        return {"pauli_string_pattern": args.pauli_pattern}
+    if args.observable_json:
+        with Path(args.observable_json).open() as f:
+            return json.load(f)
+    if hasattr(module, "build_observable"):
+        return call_builder(
+            module.build_observable,
+            nqubits=args.nqubits,
+            nlayers=args.nlayers,
+            seed=args.seed,
+        )
+    if hasattr(module, "OBSERVABLE"):
+        return module.OBSERVABLE
+    raise ValueError(
+        "No observable supplied. Define build_observable/OBSERVABLE in the case "
+        "module, or pass --pauli-pattern / --observable-json."
+    )
+
+
+def build_parallel_opts(args):
+    slicing_opts = {}
+    if args.tn_target_slices is not None:
+        slicing_opts["target_slices"] = args.tn_target_slices
+    if args.tn_target_size is not None:
+        slicing_opts["target_size"] = args.tn_target_size
+
+    opts = {
+        "slicing_opts": slicing_opts or None,
+        "search_workers": args.tn_search_workers or args.torch_threads,
+        "max_repeats": args.tn_search_repeats,
+        "max_time": args.tn_search_time,
+        "print_stats": not args.no_tn_stats,
+    }
+    if args.tn_search_backend is not None:
+        opts["search_backend"] = args.tn_search_backend
+    if args.dask_address is not None:
+        opts["dask_address"] = args.dask_address
+    if args.dask_close_workers:
+        opts["dask_close_workers"] = True
+    if args.tn_save_tree is not None:
+        opts["save_tree_path"] = args.tn_save_tree
+    if args.tn_load_tree is not None:
+        opts["load_tree_path"] = args.tn_load_tree
+    if args.tn_search_only:
+        opts["search_only"] = True
+    return opts
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Run CPU TN expectation for a custom qibo circuit module."
+    )
+    parser.add_argument("case_module", help="Python file defining build_circuit.")
+    parser.add_argument("--nqubits", type=int, required=True)
+    parser.add_argument("--nlayers", type=int, default=0)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--mpi", action="store_true")
+    parser.add_argument("--exact", action="store_true")
+    parser.add_argument("--exact-max-qubits", type=int, default=24)
+    parser.add_argument("--bond", "--bonds", dest="bond", type=optional_int, default=1024)
+    parser.add_argument("--cut-ratio", type=optional_float, default=1e-12)
+    parser.add_argument("--torch-threads", type=int, default=8)
+    parser.add_argument("--quimb-backend", choices=("numpy", "torch"), default="torch")
+    parser.add_argument("--dtype", choices=("complex128", "complex64"), default="complex128")
+    parser.add_argument("--pauli-pattern")
+    parser.add_argument("--observable-json")
+    parser.add_argument("--tn-target-slices", type=int)
+    parser.add_argument("--tn-target-size", type=int, default=2**32)
+    parser.add_argument("--tn-search-workers", type=int)
+    parser.add_argument("--tn-search-repeats", type=int, default=128)
+    parser.add_argument("--tn-search-time", type=float, default=60.0)
+    parser.add_argument("--tn-search-backend", choices=("processpool", "dask"))
+    parser.add_argument("--dask-address")
+    parser.add_argument("--dask-close-workers", action="store_true")
+    parser.add_argument("--tn-save-tree")
+    parser.add_argument("--tn-load-tree")
+    parser.add_argument("--tn-search-only", action="store_true")
+    parser.add_argument("--no-tn-stats", action="store_true")
+    args = parser.parse_args()
+
+    rank = 0
+    if args.mpi:
+        from mpi4py import MPI
+
+        rank = MPI.COMM_WORLD.Get_rank()
+
+    module = load_module(args.case_module)
+    if not hasattr(module, "build_circuit"):
+        raise ValueError("case_module must define build_circuit.")
+
+    circuit = call_builder(
+        module.build_circuit,
+        nqubits=args.nqubits,
+        nlayers=args.nlayers,
+        seed=args.seed,
+    )
+    observable = load_observable(args, module)
+
+    config = ExpectationConfig(
+        ansatz="tn",
+        mpi=args.mpi,
+        bond=args.bond,
+        cut_ratio=args.cut_ratio,
+        tensor_module="torch",
+        quimb_backend=args.quimb_backend,
+        dtype=args.dtype,
+        torch_threads=args.torch_threads,
+        parallel_opts=build_parallel_opts(args),
+    )
+
+    if rank == 0:
+        mode = "MPI" if args.mpi else "serial"
+        print(
+            f"backend=cpu ansatz=TN mode={mode} case={Path(args.case_module).name} "
+            f"nqubits={args.nqubits} nlayers={args.nlayers} seed={args.seed} "
+            f"quimb_backend={args.quimb_backend} dtype={args.dtype} "
+            f"torch_threads={args.torch_threads}",
+            flush=True,
+        )
+        print("observable exact value abs_error rel_error seconds", flush=True)
+
+    exact = None
+    if args.exact and rank == 0:
+        if args.nqubits > args.exact_max_qubits:
+            raise ValueError(
+                f"--exact is limited to {args.exact_max_qubits} qubits by default."
+            )
+        exact = exact_for_observable(circuit, observable, args.nqubits)
+
+    result = run_cpu_expectation(circuit, observable, config)
+    if args.mpi and result.rank != 0:
+        return
+
+    abs_error = float("nan") if exact is None else abs(result.value - exact)
+    rel_error = float("nan") if exact is None else abs_error / max(abs(exact), 1e-15)
+    exact_text = "nan" if exact is None else f"{exact:.16e}"
+    print(
+        f"custom {exact_text} {result.value:.16e} "
+        f"{abs_error:.6e} {rel_error:.6e} {result.seconds:.3f}",
+        flush=True,
+    )
+
+    for stat in result.parallel_stats or ():
+        cost = stat["path_cost"]
+        search_stats = stat.get("search_stats", {})
+        print(
+            "tn_term_summary "
+            f"term={stat.get('term_index', 0)} "
+            f"search_seconds={stat.get('search_seconds', float('nan')):.3f} "
+            f"contract_seconds={stat.get('contract_seconds', float('nan')):.3f} "
+            f"completed_trials={search_stats.get('completed_trials', 'na')} "
+            f"finite_trials={search_stats.get('finite_trials', 'na')} "
+            f"failed_trials={search_stats.get('failed_trials', 'na')} "
+            f"requested_trials={search_stats.get('requested_trials', 'na')} "
+            f"best_score={search_stats.get('best_score', float('nan')):.6g} "
+            f"slices={cost.get('slices')} "
+            f"log10_flops={cost.get('log10_flops', float('nan')):.3f} "
+            f"log10_write={cost.get('log10_write', float('nan')):.3f} "
+            f"log2_size={cost.get('log2_size', float('nan')):.3f} "
+            f"peak_memory_gib={cost.get('peak_memory_gib', float('nan')):.3g} "
+            f"rank_slices={stat.get('rank_slices')}",
+            flush=True,
+        )
+
+
+if __name__ == "__main__":
+    main()

tools/run_vidal_mpi_contest_cases.sh

@@ -0,0 +1,340 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Contest-style Vidal/MPI MPS cases.
+#
+# Usage:
+#   tools/run_vidal_mpi_contest_cases.sh main1
+#   tools/run_vidal_mpi_contest_cases.sh main2
+#   tools/run_vidal_mpi_contest_cases.sh strong
+#   tools/run_vidal_mpi_contest_cases.sh all
+#
+# Common overrides:
+#   PYTHON_BIN=.venv/bin/python
+#   MPIEXEC=mpiexec
+#   MPIEXEC_FULL="mpirun -np 4 -hostfile /home/yx/qibotn/hostfile -perhost 2"
+#   HOSTFILE=hostfile                 # optional; used only if the file exists
+#   RANKS=8
+#   TORCH_THREADS=8
+#   CUT_RATIO=1e-12
+#   OBS_FILTER="boundary_ZZ_q2 ring_xz dense3_spread complex_iZ0"
+#
+# Per-case overrides:
+#   MAIN1_NQ=128 MAIN1_LAYERS=50 MAIN1_BOND=1024 MAIN1_SEED=31001
+#   MAIN2_NQ=128 MAIN2_LAYERS=64 MAIN2_BOND=2048 MAIN2_SEED=31002
+#   STRONG_NQ=256 STRONG_LAYERS=64 STRONG_BOND=2048 STRONG_SEED=41001
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+cd "$ROOT_DIR"
+
+PYTHON_BIN="${PYTHON_BIN:-.venv/bin/python}"
+MPIEXEC="${MPIEXEC:-mpiexec}"
+HOSTFILE="${HOSTFILE:-}"
+RANKS="${RANKS:-4}"
+TORCH_THREADS="${TORCH_THREADS:-1}"
+CUT_RATIO="${CUT_RATIO:-1e-12}"
+OBS_FILTER="${OBS_FILTER:-}"
+
+RUNNER_DIR="$ROOT_DIR/.tmp"
+mkdir -p "$RUNNER_DIR"
+RUNNER="$(mktemp "$RUNNER_DIR/qibotn_vidal_contest.XXXXXX.py")"
+cleanup() {
+  rm -f "$RUNNER"
+}
+trap cleanup EXIT
+
+cat > "$RUNNER" <<'PY'
+from __future__ import annotations
+
+import argparse
+import math
+import time
+
+import numpy as np
+from mpi4py import MPI
+from qibo import Circuit, gates, hamiltonians
+from qibo.symbols import X, Y, Z
+
+from qibotn.backends.vidal import VidalBackend
+
+
+def set_torch_threads(nthreads):
+    try:
+        import torch
+
+        torch.set_num_threads(nthreads)
+    except Exception:
+        pass
+
+
+def build_circuit(kind, nqubits, nlayers, seed):
+    rng = np.random.default_rng(seed)
+    circuit = Circuit(nqubits)
+
+    for layer in range(nlayers):
+        for q in range(nqubits):
+            circuit.add(gates.RY(q, theta=rng.uniform(-math.pi, math.pi)))
+            circuit.add(gates.RZ(q, theta=rng.uniform(-math.pi, math.pi)))
+            if kind in ("rxx_rzz", "scramble"):
+                circuit.add(gates.RX(q, theta=rng.uniform(-math.pi, math.pi)))
+
+        if kind == "reversed_cnot":
+            for q in range(0, nqubits - 1, 2):
+                circuit.add(gates.CNOT(q + 1, q) if layer % 2 else gates.CNOT(q, q + 1))
+            for q in range(1, nqubits - 1, 2):
+                circuit.add(gates.CNOT(q + 1, q) if layer % 2 == 0 else gates.CNOT(q, q + 1))
+        elif kind == "rxx_rzz":
+            for q in range(layer % 2, nqubits - 1, 2):
+                circuit.add(gates.RXX(q, q + 1, theta=rng.uniform(-0.9, 0.9)))
+                circuit.add(gates.RZZ(q, q + 1, theta=rng.uniform(-0.9, 0.9)))
+        elif kind == "scramble":
+            for q in range(layer % 2, nqubits - 1, 2):
+                circuit.add(gates.RXX(q, q + 1, theta=rng.uniform(-0.8, 0.8)))
+                circuit.add(gates.RZZ(q, q + 1, theta=rng.uniform(-0.8, 0.8)))
+                if layer % 5 == 4:
+                    circuit.add(gates.SWAP(q, q + 1))
+        else:
+            raise ValueError(f"Unknown circuit kind {kind!r}.")
+
+    return circuit
+
+
+def ring_xz(nqubits):
+    form = 0
+    for q in range(nqubits):
+        form += 0.5 * X(q) * Z((q + 1) % nqubits)
+    return hamiltonians.SymbolicHamiltonian(form=form)
+
+
+def open_zz(nqubits):
+    form = 0
+    for q in range(nqubits - 1):
+        form += (1.0 / (nqubits - 1)) * Z(q) * Z(q + 1)
+    return hamiltonians.SymbolicHamiltonian(form=form)
+
+
+def range2_xx(nqubits):
+    form = 0
+    for q in range(nqubits - 2):
+        form += (1.0 / (nqubits - 2)) * X(q) * X(q + 2)
+    return hamiltonians.SymbolicHamiltonian(form=form)
+
+
+def dense_observable(nqubits, qubits, seed, dim):
+    rng = np.random.default_rng(seed)
+    raw = rng.normal(size=(dim, dim)) + 1j * rng.normal(size=(dim, dim))
+    matrix = (raw + raw.conj().T) / 2.0
+    matrix = matrix / np.linalg.norm(matrix)
+    return {"matrix": matrix, "qubits": list(qubits)}
+
+
+def observables_for_case(nqubits, seed):
+    q1 = nqubits // 4
+    q2 = nqubits // 2
+    q3 = (3 * nqubits) // 4
+    last = nqubits - 1
+
+    return [
+        ("boundary_ZZ_q1", hamiltonians.SymbolicHamiltonian(form=Z(q1 - 1) * Z(q1))),
+        ("boundary_ZZ_q2", hamiltonians.SymbolicHamiltonian(form=Z(q2 - 1) * Z(q2))),
+        ("boundary_ZZ_q3", hamiltonians.SymbolicHamiltonian(form=Z(q3 - 1) * Z(q3))),
+        (
+            "long_Z_5_sites",
+            hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(q1) * Z(q2) * Z(q3) * Z(last)),
+        ),
+        (
+            "mixed_XZYZX",
+            hamiltonians.SymbolicHamiltonian(form=X(0) * Z(q1) * Y(q2) * Z(q3) * X(last)),
+        ),
+        ("ring_xz", ring_xz(nqubits)),
+        ("open_zz", open_zz(nqubits)),
+        ("range2_xx", range2_xx(nqubits)),
+        ("complex_iZ0", hamiltonians.SymbolicHamiltonian(form=1.0j * Z(0))),
+        ("dense2_mid", dense_observable(nqubits, (q2 - 1, q2), seed + 101, 4)),
+        ("dense3_spread", dense_observable(nqubits, (q1, q2, q3), seed + 202, 8)),
+    ]
+
+
+def run_case(args):
+    set_torch_threads(args.torch_threads)
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    size = comm.Get_size()
+
+    circuit = build_circuit(args.kind, args.nqubits, args.nlayers, args.seed)
+    observables = observables_for_case(args.nqubits, args.seed)
+    if args.obs_filter:
+        wanted = set(args.obs_filter.split(","))
+        observables = [(name, obs) for name, obs in observables if name in wanted]
+        if not observables:
+            raise ValueError(f"OBS_FILTER matched no observables: {args.obs_filter!r}")
+
+    if rank == 0:
+        print("=" * 88, flush=True)
+        print(
+            "case "
+            f"label={args.label} kind={args.kind} ranks={size} "
+            f"nqubits={args.nqubits} nlayers={args.nlayers} gates={len(circuit.queue)} "
+            f"bond={args.bond} cut_ratio={args.cut_ratio:g} "
+            f"torch_threads={args.torch_threads} seed={args.seed} "
+            f"obs_filter={args.obs_filter or 'all'}",
+            flush=True,
+        )
+        print(
+            "observable value seconds trunc_sum trunc_max status",
+            flush=True,
+        )
+
+    for obs_name, observable in observables:
+        backend = VidalBackend()
+        backend.configure_tn_simulation(
+            max_bond_dimension=args.bond,
+            cut_ratio=args.cut_ratio,
+            tensor_module="torch",
+            mpi_approach="CT",
+            mpi_num_procs=size,
+            fallback=False,
+        )
+
+        comm.Barrier()
+        start = time.perf_counter()
+        try:
+            value = backend.expectation(
+                circuit,
+                observable,
+                preprocess=True,
+                compile_circuit=False,
+            )
+            status = "ok"
+        except Exception as exc:  # pragma: no cover - printed for manual runs
+            value = np.nan
+            status = type(exc).__name__ + ":" + str(exc).split("\n", 1)[0]
+        seconds = time.perf_counter() - start
+
+        if rank == 0:
+            print(
+                f"{obs_name} {value!r} {seconds:.3f} "
+                f"{backend.last_truncation_error:.6e} "
+                f"{backend.last_max_truncation_error:.6e} {status}",
+                flush=True,
+            )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--label", required=True)
+    parser.add_argument("--kind", choices=("reversed_cnot", "rxx_rzz", "scramble"), required=True)
+    parser.add_argument("--nqubits", type=int, required=True)
+    parser.add_argument("--nlayers", type=int, required=True)
+    parser.add_argument("--bond", type=int, required=True)
+    parser.add_argument("--cut-ratio", type=float, required=True)
+    parser.add_argument("--seed", type=int, required=True)
+    parser.add_argument("--torch-threads", type=int, required=True)
+    parser.add_argument("--obs-filter", default="")
+    run_case(parser.parse_args())
+
+
+if __name__ == "__main__":
+    main()
+PY
+
+if [[ -n "${MPIEXEC_FULL:-}" ]]; then
+  read -r -a mpi_prefix <<< "$MPIEXEC_FULL"
+else
+  mpi_prefix=("$MPIEXEC")
+  if [[ -n "$HOSTFILE" && -f "$HOSTFILE" ]]; then
+    mpi_prefix+=("-hostfile" "$HOSTFILE")
+  fi
+  mpi_prefix+=("-n" "$RANKS")
+fi
+
+run_case() {
+  local label="$1"
+  local kind="$2"
+  local nq="$3"
+  local layers="$4"
+  local bond="$5"
+  local seed="$6"
+
+  echo
+  echo "Running $label: kind=$kind nqubits=$nq layers=$layers bond=$bond seed=$seed"
+  echo "MPI: ${mpi_prefix[*]}"
+  "${mpi_prefix[@]}" "$PYTHON_BIN" -u "$ROOT_DIR/tools/vidal_mpi_contest_runner.py" \
+    --label "$label" \
+    --kind "$kind" \
+    --nqubits "$nq" \
+    --nlayers "$layers" \
+    --bond "$bond" \
+    --cut-ratio "$CUT_RATIO" \
+    --seed "$seed" \
+    --torch-threads "$TORCH_THREADS" \
+    --obs-filter "$(tr ' ' ',' <<< "$OBS_FILTER")"
+}
+
+case "${1:-help}" in
+  main1)
+    run_case \
+      "main1-reversed-cnot" \
+      "reversed_cnot" \
+      "${MAIN1_NQ:-128}" \
+      "${MAIN1_LAYERS:-50}" \
+      "${MAIN1_BOND:-1024}" \
+      "${MAIN1_SEED:-31001}"
+    ;;
+  main2)
+    run_case \
+      "main2-rxx-rzz" \
+      "rxx_rzz" \
+      "${MAIN2_NQ:-128}" \
+      "${MAIN2_LAYERS:-64}" \
+      "${MAIN2_BOND:-2048}" \
+      "${MAIN2_SEED:-31002}"
+    ;;
+  strong)
+    run_case \
+      "strong-scramble" \
+      "scramble" \
+      "${STRONG_NQ:-256}" \
+      "${STRONG_LAYERS:-64}" \
+      "${STRONG_BOND:-2048}" \
+      "${STRONG_SEED:-41001}"
+    ;;
+  all)
+    "$0" main1
+    "$0" main2
+    "$0" strong
+    ;;
+  smoke)
+    MAIN1_NQ="${MAIN1_NQ:-32}" \
+    MAIN1_LAYERS="${MAIN1_LAYERS:-6}" \
+    MAIN1_BOND="${MAIN1_BOND:-128}" \
+    "$0" main1
+    ;;
+  help|*)
+    cat >&2 <<'EOF'
+Usage: tools/run_vidal_mpi_contest_cases.sh [main1|main2|strong|all|smoke]
+
+Cases:
+  main1   128 qubits, 50 layers, reversed-CNOT brickwall, chi=1024
+  main2   128 qubits, 64 layers, RXX/RZZ brickwall, chi=2048
+  strong  256 qubits, 64 layers, RXX/RZZ + periodic SWAP scramble, chi=2048
+  smoke   Small syntax/runtime check of main1
+
+Common overrides:
+  PYTHON_BIN=.venv/bin/python
+  MPIEXEC=mpiexec
+  MPIEXEC_FULL="mpirun -np 4 -hostfile /home/yx/qibotn/hostfile -perhost 2"
+  HOSTFILE=hostfile
+  RANKS=8
+  TORCH_THREADS=8
+  CUT_RATIO=1e-12
+  OBS_FILTER="boundary_ZZ_q2 ring_xz dense3_spread complex_iZ0"
+
+Per-case overrides:
+  MAIN1_NQ=128 MAIN1_LAYERS=50 MAIN1_BOND=1024 MAIN1_SEED=31001
+  MAIN2_NQ=128 MAIN2_LAYERS=64 MAIN2_BOND=2048 MAIN2_SEED=31002
+  STRONG_NQ=256 STRONG_LAYERS=64 STRONG_BOND=2048 STRONG_SEED=41001
+EOF
+    exit 2
+    ;;
+esac

tools/slice_existing_tree.py

@@ -0,0 +1,59 @@
+"""Slice an existing saved cotengra tree without re-running path search."""
+
+from __future__ import annotations
+
+import argparse
+import pickle
+from pathlib import Path
+
+from qibotn.parallel import contraction_tree_costs
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input", help="Input pickle saved by --tn-save-tree.")
+    parser.add_argument("output", help="Output pickle path.")
+    parser.add_argument("--term", type=int, default=0)
+    parser.add_argument("--target-slices", type=int, default=2)
+    parser.add_argument("--max-repeats", type=int, default=64)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    input_path = Path(args.input)
+    output_path = Path(args.output)
+    with input_path.open("rb") as f:
+        payload = pickle.load(f)
+
+    trees = payload["trees"] if isinstance(payload, dict) else payload
+    if not isinstance(trees, (list, tuple)):
+        trees = [trees]
+    tree = trees[args.term]
+
+    print("original", contraction_tree_costs(tree), flush=True)
+    sliced = tree.slice(
+        target_slices=args.target_slices,
+        max_repeats=args.max_repeats,
+        seed=args.seed,
+    )
+    print("sliced", contraction_tree_costs(sliced), flush=True)
+    print(f"sliced_inds={sliced.sliced_inds}", flush=True)
+
+    new_trees = list(trees)
+    new_trees[args.term] = sliced
+
+    if isinstance(payload, dict):
+        out_payload = dict(payload)
+        out_payload["trees"] = new_trees
+        out_payload["costs"] = [contraction_tree_costs(t) for t in new_trees]
+        out_payload["nterms"] = len(new_trees)
+    else:
+        out_payload = new_trees
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with output_path.open("wb") as f:
+        pickle.dump(out_payload, f)
+    print(f"saved {output_path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

tools/tn_contest_runner.py

@@ -0,0 +1,435 @@
+#!/usr/bin/env python
+"""Contest-style CPU TN path search and contraction runner.
+
+This file is intentionally self-contained: define contest circuits and
+observables here, run path search once, then load the saved trees for repeated
+MPI contractions.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+import subprocess
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from urllib.parse import urlparse
+
+import numpy as np
+from qibo import Circuit, gates, hamiltonians
+from qibo.symbols import X, Y, Z
+
+ROOT = Path(__file__).resolve().parents[1]
+SRC = ROOT / "src"
+if str(SRC) not in sys.path:
+    sys.path.insert(0, str(SRC))
+
+from qibotn.expectation_runner import (  # noqa: E402
+    ExpectationConfig,
+    exact_for_observable,
+    run_cpu_expectation,
+)
+
+
+@dataclass(frozen=True)
+class CaseSpec:
+    circuit_kind: str
+    observables: tuple[str, ...]
+    nqubits: int
+    nlayers: int
+    seed: int
+    target_slices: int | None = None
+
+
+CASES = {
+    "main1": CaseSpec(
+        circuit_kind="rxx_rzz_chain",
+        observables=("ring_xz",),
+        nqubits=34,
+        nlayers=20,
+        seed=31001,
+        target_slices=None,
+    ),
+    "main2": CaseSpec(
+        circuit_kind="scramble_chain",
+        observables=("open_zz", "range2_xx"),
+        nqubits=36,
+        nlayers=18,
+        seed=31002,
+        target_slices=None,
+    ),
+    "strong": CaseSpec(
+        circuit_kind="reversed_cnot",
+        observables=("ring_xz", "long_z_string"),
+        nqubits=40,
+        nlayers=24,
+        seed=41001,
+        target_slices=None,
+    ),
+}
+
+
+def optional_int(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return int(text)
+
+
+def optional_float(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return float(text)
+
+
+def set_torch_threads(nthreads):
+    try:
+        import torch
+
+        torch.set_num_threads(nthreads)
+    except Exception:
+        pass
+
+
+def add_single_qubit_layer(circuit, nqubits, rng, include_rx=False):
+    for qubit in range(nqubits):
+        circuit.add(gates.RY(qubit, theta=rng.uniform(-math.pi, math.pi)))
+        circuit.add(gates.RZ(qubit, theta=rng.uniform(-math.pi, math.pi)))
+        if include_rx:
+            circuit.add(gates.RX(qubit, theta=rng.uniform(-math.pi, math.pi)))
+
+
+def build_circuit(kind, nqubits, nlayers, seed):
+    """Define contest circuits here."""
+    rng = np.random.default_rng(seed)
+    circuit = Circuit(nqubits)
+
+    for layer in range(nlayers):
+        if kind == "rxx_rzz_chain":
+            add_single_qubit_layer(circuit, nqubits, rng, include_rx=True)
+            for qubit in range(layer % 2, nqubits - 1, 2):
+                circuit.add(gates.RXX(qubit, qubit + 1, theta=rng.uniform(-0.9, 0.9)))
+                circuit.add(gates.RZZ(qubit, qubit + 1, theta=rng.uniform(-0.9, 0.9)))
+
+        elif kind == "scramble_chain":
+            add_single_qubit_layer(circuit, nqubits, rng, include_rx=True)
+            for qubit in range(layer % 2, nqubits - 1, 2):
+                circuit.add(gates.RXX(qubit, qubit + 1, theta=rng.uniform(-0.8, 0.8)))
+                circuit.add(gates.RZZ(qubit, qubit + 1, theta=rng.uniform(-0.8, 0.8)))
+                if layer % 5 == 4:
+                    circuit.add(gates.SWAP(qubit, qubit + 1))
+
+        elif kind == "reversed_cnot":
+            add_single_qubit_layer(circuit, nqubits, rng)
+            for qubit in range(0, nqubits - 1, 2):
+                gate = gates.CNOT(qubit + 1, qubit) if layer % 2 else gates.CNOT(qubit, qubit + 1)
+                circuit.add(gate)
+            for qubit in range(1, nqubits - 1, 2):
+                gate = gates.CNOT(qubit + 1, qubit) if layer % 2 == 0 else gates.CNOT(qubit, qubit + 1)
+                circuit.add(gate)
+
+        else:
+            raise ValueError(f"Unknown circuit kind {kind!r}.")
+
+    return circuit
+
+
+def pauli_sum_observable(kind, nqubits, seed):
+    """Define contest observables here.
+
+    TN path currently expects Pauli products / SymbolicHamiltonian terms.
+    Keep production contest observables Hermitian unless complex output is
+    explicitly required by the scoring rule.
+    """
+    del seed
+    if kind == "ring_xz":
+        form = 0
+        for qubit in range(nqubits):
+            form += 0.5 * X(qubit) * Z((qubit + 1) % nqubits)
+        return hamiltonians.SymbolicHamiltonian(form=form)
+
+    if kind == "open_zz":
+        form = 0
+        for qubit in range(nqubits - 1):
+            form += (1.0 / max(1, nqubits - 1)) * Z(qubit) * Z(qubit + 1)
+        return hamiltonians.SymbolicHamiltonian(form=form)
+
+    if kind == "range2_xx":
+        form = 0
+        for qubit in range(nqubits - 2):
+            form += (1.0 / max(1, nqubits - 2)) * X(qubit) * X(qubit + 2)
+        return hamiltonians.SymbolicHamiltonian(form=form)
+
+    if kind == "long_z_string":
+        stride = max(1, nqubits // 16)
+        form = None
+        for qubit in range(0, nqubits, stride):
+            form = Z(qubit) if form is None else form * Z(qubit)
+        return hamiltonians.SymbolicHamiltonian(form=form)
+
+    if kind == "mixed_local":
+        q1 = nqubits // 4
+        q2 = nqubits // 2
+        q3 = (3 * nqubits) // 4
+        form = 0.25 * X(0) - 0.5 * Z(nqubits - 1)
+        form += 0.125 * X(q1) * Z(q2) * Y(q3)
+        return hamiltonians.SymbolicHamiltonian(form=form)
+
+    raise ValueError(f"Unknown observable kind {kind!r}.")
+
+
+def tree_path(tree_dir, case_name, obs_name, nqubits, nlayers, target_slices):
+    slice_label = "auto" if target_slices is None else f"s{target_slices}"
+    return (
+        Path(tree_dir)
+        / f"{case_name}_{obs_name}_{nqubits}q{nlayers}l_{slice_label}.pkl"
+    )
+
+
+def build_parallel_opts(args, tree_file=None, search_only=False):
+    slicing_opts = {}
+    if args.tn_target_slices is not None:
+        slicing_opts["target_slices"] = args.tn_target_slices
+    if args.tn_target_size is not None:
+        slicing_opts["target_size"] = args.tn_target_size
+
+    opts = {
+        "slicing_opts": slicing_opts or None,
+        "search_workers": args.tn_search_workers or args.torch_threads,
+        "max_repeats": args.tn_search_repeats,
+        "max_time": args.tn_search_time,
+        "print_stats": not args.no_tn_stats,
+    }
+    if args.tn_search_backend is not None:
+        opts["search_backend"] = args.tn_search_backend
+    if args.dask_address is not None:
+        opts["dask_address"] = args.dask_address
+    if args.dask_close_workers:
+        opts["dask_close_workers"] = True
+    if args.tn_debug_trials:
+        opts["debug_trials"] = True
+    if search_only:
+        opts["search_only"] = True
+        opts["save_tree_path"] = str(tree_file)
+    elif tree_file is not None:
+        opts["load_tree_path"] = str(tree_file)
+    return opts
+
+
+def run_one(args, case_name, obs_name, mode):
+    case = CASES[case_name]
+    circuit = build_circuit(case.circuit_kind, args.nqubits, args.nlayers, args.seed)
+    observable = pauli_sum_observable(obs_name, args.nqubits, args.seed)
+    path = tree_path(
+        args.tree_dir,
+        case_name,
+        obs_name,
+        args.nqubits,
+        args.nlayers,
+        args.tn_target_slices,
+    )
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    rank = 0
+    if args.mpi:
+        from mpi4py import MPI
+
+        rank = MPI.COMM_WORLD.Get_rank()
+
+    if rank == 0:
+        print("=" * 88, flush=True)
+        print(
+            f"mode={mode} case={case_name} circuit={case.circuit_kind} "
+            f"observable={obs_name} nqubits={args.nqubits} nlayers={args.nlayers} "
+            f"seed={args.seed} gates={len(circuit.queue)} tree={path}",
+            flush=True,
+        )
+
+    if mode == "contract" and not path.exists():
+        raise FileNotFoundError(f"Missing tree file: {path}. Run search first.")
+
+    exact = None
+    if args.exact and rank == 0 and mode != "search":
+        if args.nqubits > args.exact_max_qubits:
+            raise ValueError(
+                f"--exact is limited to {args.exact_max_qubits} qubits by default."
+            )
+        exact = exact_for_observable(circuit, observable, args.nqubits)
+
+    config = ExpectationConfig(
+        ansatz="tn",
+        mpi=args.mpi,
+        bond=args.bond,
+        cut_ratio=args.cut_ratio,
+        tensor_module="torch",
+        quimb_backend=args.quimb_backend,
+        dtype=args.dtype,
+        torch_threads=args.torch_threads,
+        parallel_opts=build_parallel_opts(
+            args,
+            tree_file=path,
+            search_only=(mode == "search"),
+        ),
+    )
+    result = run_cpu_expectation(circuit, observable, config)
+    if args.mpi and result.rank != 0:
+        return
+
+    if mode == "search":
+        print(f"searched observable={obs_name} tree={path}", flush=True)
+    else:
+        abs_error = float("nan") if exact is None else abs(result.value - exact)
+        rel_error = float("nan") if exact is None else abs_error / max(abs(exact), 1e-15)
+        exact_text = "nan" if exact is None else f"{exact:.16e}"
+        print(
+            f"result observable={obs_name} exact={exact_text} "
+            f"value={result.value:.16e} abs_error={abs_error:.6e} "
+            f"rel_error={rel_error:.6e} seconds={result.seconds:.3f}",
+            flush=True,
+        )
+
+    for stat in result.parallel_stats or ():
+        cost = stat["path_cost"]
+        search_stats = stat.get("search_stats", {})
+        print(
+            "tn_term_summary "
+            f"observable={obs_name} "
+            f"term={stat.get('term_index', 0)} "
+            f"search_seconds={stat.get('search_seconds', float('nan')):.3f} "
+            f"contract_seconds={stat.get('contract_seconds', float('nan')):.3f} "
+            f"completed_trials={search_stats.get('completed_trials', 'na')} "
+            f"finite_trials={search_stats.get('finite_trials', 'na')} "
+            f"failed_trials={search_stats.get('failed_trials', 'na')} "
+            f"requested_trials={search_stats.get('requested_trials', 'na')} "
+            f"best_score={search_stats.get('best_score', float('nan')):.6g} "
+            f"slices={cost.get('slices')} "
+            f"log10_flops={cost.get('log10_flops', float('nan')):.3f} "
+            f"log10_write={cost.get('log10_write', float('nan')):.3f} "
+            f"log2_size={cost.get('log2_size', float('nan')):.3f} "
+            f"peak_memory_gib={cost.get('peak_memory_gib', float('nan')):.3g} "
+            f"rank_slices={stat.get('rank_slices')}",
+            flush=True,
+        )
+
+
+def selected_observables(args, case):
+    if args.observables:
+        return tuple(args.observables)
+    if args.obs_filter:
+        return tuple(x.strip() for x in args.obs_filter.split(",") if x.strip())
+    return case.observables
+
+
+def apply_case_defaults(args):
+    case = CASES[args.case]
+    if args.nqubits is None:
+        args.nqubits = case.nqubits
+    if args.nlayers is None:
+        args.nlayers = case.nlayers
+    if args.seed is None:
+        args.seed = case.seed
+    if args.tn_target_slices is None:
+        args.tn_target_slices = case.target_slices
+    args.observables = selected_observables(args, case)
+
+
+def stop_dask_cluster(args):
+    if args.keep_dask or args.tn_search_backend != "dask" or not args.dask_address:
+        return
+    script = ROOT / "tools" / "manage_tn_dask_cluster.sh"
+    if not script.exists():
+        print(f"dask_stop_skipped reason=missing_script path={script}", flush=True)
+        return
+
+    env = os.environ.copy()
+    parsed = urlparse(args.dask_address)
+    if parsed.hostname:
+        env.setdefault("SCHEDULER_HOST", parsed.hostname)
+    if parsed.port:
+        env.setdefault("SCHEDULER_PORT", str(parsed.port))
+
+    print("dask_stop_after_search start", flush=True)
+    subprocess.run([str(script), "stop"], cwd=str(ROOT), env=env, check=False)
+    print("dask_stop_after_search done", flush=True)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("mode", choices=("search", "contract", "all", "validate", "list"))
+    parser.add_argument("--case", choices=sorted(CASES), default="main1")
+    parser.add_argument("--observables", nargs="+")
+    parser.add_argument("--obs-filter", default="")
+    parser.add_argument("--tree-dir", default="trees/contest_tn")
+    parser.add_argument("--nqubits", type=int)
+    parser.add_argument("--nlayers", type=int)
+    parser.add_argument("--seed", type=int)
+    parser.add_argument("--mpi", action="store_true")
+    parser.add_argument("--exact", action="store_true")
+    parser.add_argument("--exact-max-qubits", type=int, default=24)
+    parser.add_argument("--bond", "--bonds", dest="bond", type=optional_int, default=1024)
+    parser.add_argument("--cut-ratio", type=optional_float, default=1e-12)
+    parser.add_argument("--torch-threads", type=int, default=8)
+    parser.add_argument("--quimb-backend", choices=("numpy", "torch"), default="torch")
+    parser.add_argument("--dtype", choices=("complex128", "complex64"), default="complex64")
+    parser.add_argument("--tn-target-slices", type=int)
+    parser.add_argument("--tn-target-size", type=int, default=2**32)
+    parser.add_argument("--tn-search-workers", type=int)
+    parser.add_argument("--tn-search-repeats", type=int, default=2048)
+    parser.add_argument("--tn-search-time", type=float, default=300.0)
+    parser.add_argument(
+        "--tn-search-backend",
+        choices=("processpool", "dask"),
+        default="dask",
+        help=(
+            "Path-search backend. Defaults to dask. Without --dask-address, "
+            "non-MPI search starts a local dask cluster."
+        ),
+    )
+    parser.add_argument("--dask-address")
+    parser.add_argument("--dask-close-workers", action="store_true")
+    parser.add_argument(
+        "--keep-dask",
+        action="store_true",
+        help=(
+            "Keep an external dask cluster running after search. By default, "
+            "tools/manage_tn_dask_cluster.sh stop is called after search when "
+            "--dask-address is used."
+        ),
+    )
+    parser.add_argument(
+        "--tn-debug-trials",
+        action="store_true",
+        help="Print dask worker summary and per-trial start/done logs.",
+    )
+    parser.add_argument("--no-tn-stats", action="store_true")
+    args = parser.parse_args()
+
+    if args.mode == "list":
+        for name, case in CASES.items():
+            print(
+                f"{name}: circuit={case.circuit_kind} "
+                f"observables={','.join(case.observables)} "
+                f"nqubits={case.nqubits} nlayers={case.nlayers} "
+                f"seed={case.seed} target_slices={case.target_slices}"
+            )
+        return
+
+    apply_case_defaults(args)
+    set_torch_threads(args.torch_threads)
+
+    modes = ("search", "contract") if args.mode == "all" else (args.mode,)
+    if args.mode == "validate":
+        args.exact = True
+        args.nqubits = min(args.nqubits, args.exact_max_qubits)
+        modes = ("search", "contract")
+
+    for mode in modes:
+        for obs_name in args.observables:
+            run_one(args, args.case, obs_name, mode)
+        if mode == "search":
+            stop_dask_cluster(args)
+
+
+if __name__ == "__main__":
+    main()

tools/torch_profile_tn_complex64.py

@@ -0,0 +1,114 @@
+"""Run the 34q/20L TN complex64 benchmark under torch.profiler briefly."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import signal
+import sys
+from pathlib import Path
+
+from mpi4py import MPI
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--seconds", type=float, default=30.0)
+    parser.add_argument("--out-dir", default="torch_profiles/tn_complex64")
+    parser.add_argument("--torch-threads", type=int, default=48)
+    args = parser.parse_args()
+
+    repo_root = Path(__file__).resolve().parents[1]
+    os.chdir(repo_root)
+    sys.path.insert(0, str(repo_root))
+
+    import torch
+    from torch.profiler import ProfilerActivity, profile
+
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    size = comm.Get_size()
+    out_dir = Path(args.out_dir)
+    if rank == 0:
+        out_dir.mkdir(parents=True, exist_ok=True)
+    comm.Barrier()
+
+    torch.set_num_threads(args.torch_threads)
+
+    def run_benchmark():
+        import benchmark_cpu_expectation
+
+        sys.argv = [
+            "benchmark_cpu_expectation.py",
+            "--mpi",
+            "--ansatz",
+            "tn",
+            "--nqubits",
+            "34",
+            "--nlayers",
+            "20",
+            "--circuits",
+            "rxx_rzz",
+            "--pauli-pattern",
+            "XZ",
+            "--tn-load-tree",
+            "trees/rxx_rzz_34q20l_s4.pkl",
+            "--quimb-backend",
+            "torch",
+            "--torch-threads",
+            str(args.torch_threads),
+            "--dtype",
+            "complex64",
+        ]
+        benchmark_cpu_expectation.main()
+
+    trace_path = out_dir / f"rank{rank}_trace.json"
+    stacks_path = out_dir / f"rank{rank}_stacks.txt"
+    summary_path = out_dir / f"rank{rank}_summary.txt"
+
+    prof = profile(
+        activities=[ProfilerActivity.CPU],
+        record_shapes=True,
+        profile_memory=True,
+        with_stack=True,
+    )
+
+    class ProfileTimeout(Exception):
+        pass
+
+    def alarm_handler(signum, frame):
+        raise ProfileTimeout()
+
+    old_handler = signal.signal(signal.SIGALRM, alarm_handler)
+    signal.setitimer(signal.ITIMER_REAL, args.seconds)
+    try:
+        with prof:
+            try:
+                run_benchmark()
+            except ProfileTimeout:
+                pass
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+        signal.signal(signal.SIGALRM, old_handler)
+
+    prof.export_chrome_trace(str(trace_path))
+    try:
+        prof.export_stacks(str(stacks_path), "self_cpu_time_total")
+    except Exception as exc:  # pragma: no cover - diagnostic only
+        stacks_path.write_text(f"export_stacks failed: {exc}\n", encoding="utf-8")
+
+    summary = prof.key_averages(group_by_stack_n=5).table(
+        sort_by="self_cpu_time_total",
+        row_limit=40,
+    )
+    summary_path.write_text(summary, encoding="utf-8")
+
+    print(
+        f"torch_profile_done rank={rank}/{size} "
+        f"trace={trace_path} summary={summary_path}",
+        flush=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

tools/vidal_mpi_contest_runner.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+import argparse
+import math
+import time
+
+import numpy as np
+from mpi4py import MPI
+from qibo import Circuit, gates, hamiltonians
+from qibo.symbols import X, Y, Z
+
+from qibotn.backends.vidal import VidalBackend
+
+
+def optional_int(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return int(text)
+
+
+def optional_float(text):
+    if isinstance(text, str) and text.lower() in {"none", "null", "inf", "unlimited"}:
+        return None
+    return float(text)
+
+
+def format_optional(value, fmt="g"):
+    return "None" if value is None else format(value, fmt)
+
+
+def set_torch_threads(nthreads):
+    try:
+        import torch
+
+        torch.set_num_threads(nthreads)
+    except Exception:
+        pass
+
+
+def build_circuit(kind, nqubits, nlayers, seed):
+    rng = np.random.default_rng(seed)
+    circuit = Circuit(nqubits)
+
+    for layer in range(nlayers):
+        for q in range(nqubits):
+            circuit.add(gates.RY(q, theta=rng.uniform(-math.pi, math.pi)))
+            circuit.add(gates.RZ(q, theta=rng.uniform(-math.pi, math.pi)))
+            if kind in ("rxx_rzz", "scramble"):
+                circuit.add(gates.RX(q, theta=rng.uniform(-math.pi, math.pi)))
+
+        if kind == "reversed_cnot":
+            for q in range(0, nqubits - 1, 2):
+                circuit.add(gates.CNOT(q + 1, q) if layer % 2 else gates.CNOT(q, q + 1))
+            for q in range(1, nqubits - 1, 2):
+                circuit.add(gates.CNOT(q + 1, q) if layer % 2 == 0 else gates.CNOT(q, q + 1))
+        elif kind == "rxx_rzz":
+            for q in range(layer % 2, nqubits - 1, 2):
+                circuit.add(gates.RXX(q, q + 1, theta=rng.uniform(-0.9, 0.9)))
+                circuit.add(gates.RZZ(q, q + 1, theta=rng.uniform(-0.9, 0.9)))
+        elif kind == "scramble":
+            for q in range(layer % 2, nqubits - 1, 2):
+                circuit.add(gates.RXX(q, q + 1, theta=rng.uniform(-0.8, 0.8)))
+                circuit.add(gates.RZZ(q, q + 1, theta=rng.uniform(-0.8, 0.8)))
+                if layer % 5 == 4:
+                    circuit.add(gates.SWAP(q, q + 1))
+        else:
+            raise ValueError(f"Unknown circuit kind {kind!r}.")
+
+    return circuit
+
+
+def ring_xz(nqubits):
+    form = 0
+    for q in range(nqubits):
+        form += 0.5 * X(q) * Z((q + 1) % nqubits)
+    return hamiltonians.SymbolicHamiltonian(form=form)
+
+
+def open_zz(nqubits):
+    form = 0
+    for q in range(nqubits - 1):
+        form += (1.0 / (nqubits - 1)) * Z(q) * Z(q + 1)
+    return hamiltonians.SymbolicHamiltonian(form=form)
+
+
+def range2_xx(nqubits):
+    form = 0
+    for q in range(nqubits - 2):
+        form += (1.0 / (nqubits - 2)) * X(q) * X(q + 2)
+    return hamiltonians.SymbolicHamiltonian(form=form)
+
+
+def dense_observable(nqubits, qubits, seed, dim):
+    rng = np.random.default_rng(seed)
+    raw = rng.normal(size=(dim, dim)) + 1j * rng.normal(size=(dim, dim))
+    matrix = (raw + raw.conj().T) / 2.0
+    matrix = matrix / np.linalg.norm(matrix)
+    return {"matrix": matrix, "qubits": list(qubits)}
+
+
+def observables_for_case(nqubits, seed):
+    q1 = nqubits // 4
+    q2 = nqubits // 2
+    q3 = (3 * nqubits) // 4
+    last = nqubits - 1
+
+    return [
+        ("boundary_ZZ_q1", hamiltonians.SymbolicHamiltonian(form=Z(q1 - 1) * Z(q1))),
+        ("boundary_ZZ_q2", hamiltonians.SymbolicHamiltonian(form=Z(q2 - 1) * Z(q2))),
+        ("boundary_ZZ_q3", hamiltonians.SymbolicHamiltonian(form=Z(q3 - 1) * Z(q3))),
+        (
+            "long_Z_5_sites",
+            hamiltonians.SymbolicHamiltonian(form=Z(0) * Z(q1) * Z(q2) * Z(q3) * Z(last)),
+        ),
+        (
+            "mixed_XZYZX",
+            hamiltonians.SymbolicHamiltonian(form=X(0) * Z(q1) * Y(q2) * Z(q3) * X(last)),
+        ),
+        ("ring_xz", ring_xz(nqubits)),
+        ("open_zz", open_zz(nqubits)),
+        ("range2_xx", range2_xx(nqubits)),
+        ("complex_iZ0", hamiltonians.SymbolicHamiltonian(form=1.0j * Z(0))),
+        ("dense2_mid", dense_observable(nqubits, (q2 - 1, q2), seed + 101, 4)),
+        ("dense3_spread", dense_observable(nqubits, (q1, q2, q3), seed + 202, 8)),
+    ]
+
+
+def run_case(args):
+    set_torch_threads(args.torch_threads)
+    comm = MPI.COMM_WORLD
+    rank = comm.Get_rank()
+    size = comm.Get_size()
+
+    circuit = build_circuit(args.kind, args.nqubits, args.nlayers, args.seed)
+    observables = observables_for_case(args.nqubits, args.seed)
+    if args.obs_filter:
+        wanted = set(args.obs_filter.split(","))
+        observables = [(name, obs) for name, obs in observables if name in wanted]
+        if not observables:
+            raise ValueError(f"OBS_FILTER matched no observables: {args.obs_filter!r}")
+
+    if rank == 0:
+        print("=" * 88, flush=True)
+        print(
+            "case "
+            f"label={args.label} kind={args.kind} ranks={size} "
+            f"nqubits={args.nqubits} nlayers={args.nlayers} gates={len(circuit.queue)} "
+            f"bond={format_optional(args.bond)} "
+            f"cut_ratio={format_optional(args.cut_ratio)} "
+            f"torch_threads={args.torch_threads} seed={args.seed} "
+            f"obs_filter={args.obs_filter or 'all'}",
+            flush=True,
+        )
+        print(
+            "observable value seconds trunc_sum trunc_max status",
+            flush=True,
+        )
+
+    for obs_name, observable in observables:
+        backend = VidalBackend()
+        backend.configure_tn_simulation(
+            max_bond_dimension=args.bond,
+            cut_ratio=args.cut_ratio,
+            tensor_module="torch",
+            mpi_approach="CT",
+            mpi_num_procs=size,
+            fallback=False,
+        )
+
+        comm.Barrier()
+        start = time.perf_counter()
+        try:
+            value = backend.expectation(
+                circuit,
+                observable,
+                preprocess=True,
+                compile_circuit=False,
+            )
+            status = "ok"
+        except Exception as exc:  # pragma: no cover - printed for manual runs
+            value = np.nan
+            status = type(exc).__name__ + ":" + str(exc).split("\n", 1)[0]
+        seconds = time.perf_counter() - start
+
+        if rank == 0:
+            print(
+                f"{obs_name} {value!r} {seconds:.3f} "
+                f"{backend.last_truncation_error:.6e} "
+                f"{backend.last_max_truncation_error:.6e} {status}",
+                flush=True,
+            )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--label", required=True)
+    parser.add_argument("--kind", choices=("reversed_cnot", "rxx_rzz", "scramble"), required=True)
+    parser.add_argument("--nqubits", type=int, required=True)
+    parser.add_argument("--nlayers", type=int, required=True)
+    parser.add_argument("--bond", type=optional_int, required=True)
+    parser.add_argument("--cut-ratio", type=optional_float, required=True)
+    parser.add_argument("--seed", type=int, required=True)
+    parser.add_argument("--torch-threads", type=int, required=True)
+    parser.add_argument("--obs-filter", default="")
+    run_case(parser.parse_args())
+
+
+if __name__ == "__main__":
+    main()