dense-evolution 8.3.0__py3-none-win_amd64.whl

dense_evolution/backends/mps.py

@@ -0,0 +1,1569 @@
+"""
+MPSSimulator - Matrix Product State statevector simulator, JAX-backed.
+
+Ported from the "TurboQuant TUREQ MPSSimulator v8.2 MatryoshkaFlash"
+prototype (private research notebook, never published as part of the
+dense-evolution package). Two real bugs were found and fixed by
+independent verification against DenseSVSimulator before this module
+existed in its current form:
+
+1. The original applied Lloyd-Max quantization to the SVD singular
+   values on every truncation ("PolarQuantizer"). Measured a real ~0.5%
+   Total Variation Distance error against DenseSVSimulator on an
+   8-qubit entangling test circuit, with ZERO bond-dimension savings to
+   show for it. Dropped entirely -- this module keeps only the plain
+   adaptive SVD truncation (JSD-budget-driven bond dimension, the
+   author's own stopping criterion -- standard SVD truncation,
+   non-standard stopping metric).
+
+2. get_top_k_probable_states (originally "_extract_top_k_paths") picked
+   a single "best" bond index via argmax at each step instead of
+   correctly summing over the bond dimension. Measured 0/8 correct
+   states against the exact contraction on the same test circuit,
+   values off by ~30x. Fixed by propagating the true partial-contraction
+   vector through each bond (matches exactly, to machine precision, on
+   every state it finds) -- but note it's a genuine greedy beam search,
+   not an exact top-k finder: recall of the true top states grows with
+   beam width k but isn't guaranteed complete for any fixed k.
+
+Originally ported in plain numpy (matching the prototype), then
+converted to jax.numpy so the core tensor contractions (einsum, SVD) run
+on the same backend as the rest of dense_evolution instead of a second,
+inconsistent numerics stack. Re-verified against DenseSVSimulator after
+the conversion -- see test_mps.py.
+
+Uses whatever jax_enable_x64 precision is currently active in the
+process (does not toggle it itself) -- same convention as
+DenseSVSimulator/Chunk, which rely on the caller (dashboard_core.py's
+run_simulation) to set precision, since jax_enable_x64 is a process-wide
+flag and toggling it locally would leak to unrelated code running later
+in the same process.
+
+For circuits with LOW entanglement (product states, GHZ/Bell-like chains,
+shallow local circuits), the bond dimension stays small regardless of
+qubit count, so this scales to hundreds of qubits where DenseSVSimulator
+(or Chunk) cannot -- see get_probabilities_sampled and
+get_top_k_probable_states, neither of which ever materializes a
+(2**n,)-shaped array. For HIGHLY entangled circuits the bond dimension
+grows and this degrades back toward the same exponential cost
+DenseSVSimulator has -- it is not a universal replacement, it is
+complementary.
+"""
+
+import dataclasses
+import heapq
+import warnings
+from functools import partial
+from typing import List, Optional, Tuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from ..config import ensure_x64
+from ..physics.observables import _normalize_terms
+
+
+def _real_dtype_for(complex_dtype) -> jnp.dtype:
+    """The real dtype that pairs with a given complex dtype -- float64
+    for complex128, float32 for complex64. Diagnostic/bookkeeping arrays
+    (chi, jsd, truncation error, entanglement entropy, the compiled op
+    table) are real-valued and must follow whichever precision the
+    quantum state itself is running at, instead of a dtype literal that
+    JAX silently downcasts (with a warning) whenever x64 isn't active."""
+    return jnp.float64 if complex_dtype == jnp.complex128 else jnp.float32
+
+
+def _jsd_vectors_jax(p: jnp.ndarray, q: jnp.ndarray) -> jnp.ndarray:
+    """Same Jensen-Shannon Distance as _jsd_vectors, but returns a jnp
+    scalar instead of a Python float -- the float() cast in _jsd_vectors
+    forces concretization, which is fine called eagerly but raises
+    ConcretizationTypeError under jax.vmap/jax.jit tracing (needed by
+    _vectorized_chi_search below). Same math, no behavior difference for
+    eager callers, who go through _jsd_vectors instead."""
+    eps = 1e-12
+    p_norm = p / (jnp.sum(p) + eps)
+    q_norm = q / (jnp.sum(q) + eps)
+    m = 0.5 * (p_norm + q_norm)
+
+    def _kl(a, b):
+        mask = a > eps
+        # jnp.where still traces both branches, but log(0) only ever
+        # feeds into a term multiplied by 0 through the mask -- safe,
+        # and avoids a Python-side branch on a traced value.
+        safe_log = jnp.where(mask, jnp.log2(jnp.where(mask, a, 1.0) / (b + eps)), 0.0)
+        return jnp.sum(jnp.where(mask, a * safe_log, 0.0))
+
+    js = 0.5 * (_kl(p_norm, m) + _kl(q_norm, m))
+    # p.shape[0] is a static Python int (shape, not a value) even for a
+    # traced array under vmap, so this branch is safe at trace time.
+    dim_factor = float(np.log10(p.shape[0]) / 2.0) if p.shape[0] > 1 else 1.0
+    return jnp.sqrt(jnp.clip(js * dim_factor, 0.0, 1.0))
+
+
+def _jsd_vectors(p: jnp.ndarray, q: jnp.ndarray) -> float:
+    """Adaptive Jensen-Shannon Distance used to size the truncated bond
+    dimension: scales with log10(dim) so the same JSD budget stays
+    meaningful whether the local Hilbert space is small or large."""
+    return float(_jsd_vectors_jax(p, q))
+
+
+def _vectorized_chi_search_jax(S: jnp.ndarray, eps: float, jsd_budget: float, max_bond: int):
+    """JIT/vmap-compatible core of _vectorized_chi_search -- no
+    float()/int()/bool() concretizing casts, returns (chi_new, jsd_val) as
+    traced jnp scalars instead of Python int/float, so it's usable from
+    inside a jax.lax.scan/jax.jit region (needed by _build_mps_runner's
+    step function). Same algorithm _vectorized_chi_search wraps and casts
+    to Python types for eager callers -- see that docstring for the
+    verification this replicates the real while loop exactly."""
+    n = S.shape[0]
+    max_possible = min(n, max_bond)
+    mask_above_eps = S > eps
+    initial_chi = jnp.clip(jnp.sum(mask_above_eps), 1, max_possible)
+
+    norm_full = jnp.sum(S ** 2) + 1e-15
+    p_full = (S ** 2) / norm_full
+    # Real bug (found via a real-machine bisection down to 5e34a1e, not
+    # guessed): under jax.jit, _jsd_vectors_jax's own internal eps=1e-12
+    # mask does not reliably zero JSD contributions from p_full's
+    # already-negligible tail (~1e-18, from squaring SVD noise-floor
+    # singular values ~1e-9) when batched via jax.vmap across multiple
+    # candidates at once -- eager execution of the exact same formula on
+    # the exact same array gives the correct (zero) JSD there, jit does
+    # not. Zeroing p_full's own negligible tail here, before it ever
+    # reaches the vmapped JSD computation, verified to fix this under
+    # jit (confirmed both eager and jit now agree) without changing any
+    # value large enough to matter physically.
+    p_full = jnp.where(p_full > 1e-12, p_full, 0.0)
+
+    candidates = jnp.arange(1, max_possible + 1)
+    idx = jnp.arange(n)
+    trunc_mask = idx[None, :] < candidates[:, None]
+    p_trunc_all = jnp.where(trunc_mask, p_full[None, :], 0.0)
+
+    jsd_all = jax.vmap(lambda pt: _jsd_vectors_jax(p_full, pt))(p_trunc_all)
+
+    valid_candidate = candidates >= initial_chi
+    satisfies = (jsd_all <= jsd_budget) & valid_candidate
+    any_satisfies = jnp.any(satisfies)
+    first_idx = jnp.argmax(satisfies)
+    chi_new = jnp.where(any_satisfies, candidates[first_idx], max_possible)
+    jsd_val = jnp.where(any_satisfies, jsd_all[first_idx], jsd_all[max_possible - 1])
+    return chi_new, jsd_val
+
+
+def _vectorized_chi_search(S: jnp.ndarray, eps: float, jsd_budget: float, max_bond: int):
+    """JIT/vmap-compatible replacement for the Python `while` loop in
+    _svd_truncate (host-syncs every iteration today). Computes the JSD for
+    every candidate truncation size in one vectorized pass instead of
+    incrementing chi_new one at a time, then picks the smallest candidate
+    >= the eps-cutoff-based starting point that satisfies jsd_budget --
+    exactly replicating the while loop's search direction and stopping
+    condition (verified: 0 mismatches in chi_new/jsd_val against the real
+    while loop across 171 real _svd_truncate calls from actual entangled
+    circuits at 3 different (n_qubits, max_bond, jsd_budget) configurations,
+    including the budget-violation fallback path).
+
+    Does NOT assume JSD is monotonic in the candidate size -- it restricts
+    the search to candidates >= the same starting point the while loop
+    starts from and takes the first (smallest) one satisfying the budget,
+    same as the loop would find by incrementing, whether or not JSD happens
+    to be monotonic there.
+
+    Returns (chi_new, jsd_val) as Python int/float -- eager convenience
+    wrapper around _vectorized_chi_search_jax, for callers outside a jit
+    region (e.g. these same tests).
+    """
+    chi_new, jsd_val = _vectorized_chi_search_jax(S, eps, jsd_budget, max_bond)
+    return int(chi_new), float(jsd_val)
+
+
+def _expand_nonlocal_2q_positions(q1: int, q2: int):
+    """Pure-Python prediction of the exact adjacent-pair call sequence
+    MPSSimulator._apply_nonlocal_2q's runtime SWAP-chain dispatch produces
+    for a given (q1, q2) -- moves that decision from runtime (inside the
+    eager gate-application loop) to pre-compile time, so it can be used to
+    build a flat op list before entering a JIT region, mirroring how
+    chunk.py's _compile_multi_chunk_ops does all gate-name branching
+    outside the traced kernel.
+
+    Verified exhaustively against the real runtime dispatch: 0 mismatches
+    across all (q1, q2) pairs for n_qubits in {6, 10, 15} (330 pairs).
+
+    Returns (seq, gate_index, needs_transpose):
+      seq            -- list of (a, a+1) adjacent physical-position pairs,
+                         in the exact order _apply_nonlocal_2q would call
+                         apply_gate_2q on them.
+      gate_index      -- which position in `seq` is the REAL gate (every
+                         other position is a SWAP).
+      needs_transpose -- True iff the original q1 > q2, matching
+                         _apply_nonlocal_2q's own q1>q2 normalization
+                         (the gate tensor must be transposed via axes
+                         (1, 0, 3, 2) before use at `seq[gate_index]`,
+                         same as _apply_nonlocal_2q already does).
+    """
+    needs_transpose = q1 > q2
+    target, far = (q2, q1) if needs_transpose else (q1, q2)
+    seq = []
+    for q in range(far - 1, target, -1):
+        seq.append((q, q + 1))
+    gate_index = len(seq)
+    seq.append((target, target + 1))
+    for q in range(target + 1, far):
+        seq.append((q, q + 1))
+    return seq, gate_index, needs_transpose
+
+
+# ── gate-ID -> matrix construction, jax.lax.switch-based (JIT-fusion) ────
+# Same gate-ID table as dense_evolution/gates.py::GATE_IDS and
+# compiler.py::_apply_gate_fast_step's g_1q switch -- same formulas,
+# verified to match exactly (0 error, both eagerly and under jax.vmap/jit)
+# before being committed here. Unlike compiler.py's statevector-bitmask
+# do_2q, MPS's own gate application is an explicit tensor contraction
+# (jnp.einsum in apply_gate_2q), so the 2-qubit switch below returns the
+# full (2,2,2,2) unitary tensor directly -- same matrices mps.py's own
+# apply_cx/apply_cz/apply_swap already build by hand, generalized to a
+# traced g_id instead of one hardcoded gate per method. id=23 (SWAP) is
+# reserved-but-unused in gates.py's own table (comment there: "never
+# dispatched here, QuantumTranspiler always decomposes swap into 3xCX
+# first") -- MPS's own kernel claims it for real, since decomposing SWAP
+# into 3 CX would mean 3 SVD truncations per swap instead of 1, defeating
+# the point of keeping SWAP as its own gate for the non-adjacent-gate
+# chain (see _expand_nonlocal_2q_positions above).
+
+def _mps_1q_matrix(g_id: jnp.ndarray, param: jnp.ndarray, dtype) -> jnp.ndarray:
+    """Traced gate-ID -> (2,2) unitary, ids 0-13 (I,H,X,Y,Z,S,Sdg,T,Tdg,
+    Rx,Ry,Rz,P,SX) -- same formulas as compiler.py's g_1q switch."""
+    inv2 = jnp.asarray(1.0 / jnp.sqrt(2.0), dtype=dtype)
+    half_p = param * 0.5
+    cos_p = jnp.cos(half_p).astype(dtype)
+    sin_p = jnp.sin(half_p).astype(dtype)
+    exp_pos = jnp.exp(1j * param).astype(dtype)
+    exp_ph4 = jnp.exp(1j * jnp.pi / 4.0).astype(dtype)
+    exp_mh4 = jnp.exp(-1j * jnp.pi / 4.0).astype(dtype)
+    safe_gid = jnp.clip(g_id, 0, 13)
+    return jax.lax.switch(
+        safe_gid,
+        [
+            lambda _: jnp.eye(2, dtype=dtype),                                          # 0  I
+            lambda _: jnp.array([[inv2, inv2], [inv2, -inv2]], dtype=dtype),             # 1  H
+            lambda _: jnp.array([[0.0 + 0j, 1.0 + 0j], [1.0 + 0j, 0.0 + 0j]], dtype=dtype),  # 2  X
+            lambda _: jnp.array([[0.0 + 0j, -1j], [1j, 0.0 + 0j]], dtype=dtype),         # 3  Y
+            lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, -1.0 + 0j]], dtype=dtype),  # 4  Z
+            lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, 1j]], dtype=dtype),    # 5  S
+            lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, -1j]], dtype=dtype),   # 6  Sdg
+            lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_ph4]], dtype=dtype),  # 7  T
+            lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_mh4]], dtype=dtype),  # 8  Tdg
+            lambda _: jnp.array([[cos_p, -1j * sin_p], [-1j * sin_p, cos_p]], dtype=dtype),  # 9  Rx
+            lambda _: jnp.array([[cos_p, -sin_p], [sin_p, cos_p]], dtype=dtype),         # 10 Ry
+            lambda _: jnp.array([[jnp.exp(-1j * half_p), 0.0 + 0j],
+                                  [0.0 + 0j, jnp.exp(1j * half_p)]], dtype=dtype),        # 11 Rz
+            lambda _: jnp.array([[1.0 + 0j, 0.0 + 0j], [0.0 + 0j, exp_pos]], dtype=dtype),  # 12 P
+            lambda _: jnp.array([[0.5 + 0.5j, 0.5 - 0.5j], [0.5 - 0.5j, 0.5 + 0.5j]], dtype=dtype),  # 13 SX
+        ],
+        operand=None,
+    )
+
+
+def _mps_2q_matrix(g_id: jnp.ndarray, param: jnp.ndarray, dtype) -> jnp.ndarray:
+    """Traced gate-ID -> (2,2,2,2) unitary tensor, ids 20-25
+    (CX,CZ,CP,SWAP,CY,CRZ) -- same target gates as compiler.py's do_2q,
+    built as explicit tensors (not the statevector-bitmask approach) to
+    match MPS's own einsum-based 2-qubit gate application."""
+    exp_pos = jnp.exp(1j * param).astype(dtype)
+    exp_neg_half = jnp.exp(-1j * param * 0.5).astype(dtype)
+    exp_pos_half = jnp.exp(1j * param * 0.5).astype(dtype)
+    one = jnp.asarray(1.0 + 0j, dtype=dtype)
+    cp_diag = jnp.stack([one, one, one, exp_pos])
+    crz_diag = jnp.stack([one, one, exp_neg_half, exp_pos_half])
+    safe_idx = jnp.clip(g_id - 20, 0, 5)
+    mat = jax.lax.switch(
+        safe_idx,
+        [
+            lambda _: jnp.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], dtype=dtype),  # 20 CX
+            lambda _: jnp.diag(jnp.array([1, 1, 1, -1], dtype=dtype)),                    # 21 CZ
+            lambda _: jnp.diag(cp_diag),                                                   # 22 CP
+            lambda _: jnp.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=dtype),  # 23 SWAP
+            lambda _: jnp.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, -1j], [0, 0, 1j, 0]], dtype=dtype),  # 24 CY
+            lambda _: jnp.diag(crz_diag),                                                  # 25 CRZ
+        ],
+        operand=None,
+    )
+    return mat.reshape(2, 2, 2, 2)
+
+
+def _pad_gamma(g: jnp.ndarray, max_bond: int) -> jnp.ndarray:
+    """Zero-pads a (chiL, 2, chiR) gamma tensor up to (max_bond, 2,
+    max_bond) -- the fixed shape run_circuit_jit's scan carry needs.
+    Mathematically transparent: every contraction used here (einsum) sums
+    over the padded axis, and zero entries there contribute exactly zero,
+    verified directly (max singular-value error at machine epsilon,
+    final-state fidelity 1.0 to 1e-13, against the real eager simulator
+    across multiple (n_qubits, max_bond, jsd_budget) configurations
+    including budget-violation cases)."""
+    chi_l, d, chi_r = g.shape
+    out = jnp.zeros((max_bond, d, max_bond), dtype=g.dtype)
+    return out.at[:chi_l, :, :chi_r].set(g)
+
+
+def _pad_lambda(lam: jnp.ndarray, max_bond: int) -> jnp.ndarray:
+    """Zero-pads a (chi,) lambda vector up to (max_bond,)."""
+    n = lam.shape[0]
+    out = jnp.zeros((max_bond,), dtype=lam.dtype)
+    return out.at[:n].set(lam)
+
+
+@partial(jax.jit, static_argnums=(1, 2))
+def _pad_all_gammas(gammas, max_bond: int, dtype) -> jnp.ndarray:
+    """Batches the whole per-gamma padding+stacking loop into one
+    compiled call instead of n_qubits separate eager _pad_gamma
+    dispatches -- profiling run_circuit_jit directly found this eager
+    loop (50-51 separate JAX dispatches for N=50) cost 52-55% of total
+    run_circuit_jit time, present equally on the fuse_gates=True and
+    default paths, diluting the real speedup ratio between them.
+    Persistent-compiled-cache wrapper, keyed by the tuple of gamma
+    shapes + max_bond + dtype -- valid across the whole process
+    lifetime, same principle as _jit_mps_1q_matrix/_jit_mps_2q_matrix."""
+    return jnp.stack([_pad_gamma(g, max_bond).astype(dtype) for g in gammas])
+
+
+@partial(jax.jit, static_argnums=(1, 2))
+def _pad_all_lambdas(lambdas, max_bond: int, dtype) -> jnp.ndarray:
+    """Same persistent-compiled-cache fix as _pad_all_gammas, for
+    lambdas."""
+    return jnp.stack([_pad_lambda(l, max_bond).astype(dtype) for l in lambdas])
+
+
+def _compile_mps_ops(ops, n_qubits: int) -> List[List[float]]:
+    """Pre-compile-time (pure Python, outside any jit region) translation
+    of a circuit -- list of (name, *args) tuples/lists, same convention as
+    DenseSVSimulator.run_circuit_jit_beast_mode -- into a flat list of
+    [g_id, q1, q2, param, transpose_flag] rows for run_circuit_jit's
+    jax.lax.scan kernel. Two things happen here, deliberately outside the
+    jit region (same principle as chunk.py's _compile_multi_chunk_ops --
+    all Python-level branching on gate identity happens before tracing
+    starts, never inside the traced step function):
+
+      1. GATE_IDS (dense_evolution/gates.py) name -> numeric id lookup.
+      2. Non-adjacent 2-qubit gates, and adjacent ones with q1 > q2,
+         expanded/normalized via _expand_nonlocal_2q_positions (verified
+         above) into their SWAP-chain-equivalent sequence of
+         adjacent-ascending rows -- moves what _apply_nonlocal_2q's
+         runtime dispatch does today to pre-compile time.
+
+    id 23 (SWAP) is used for real here (not decomposed into 3xCX like
+    QuantumTranspiler.transpile would) -- see _mps_2q_matrix's docstring
+    for why: decomposing would mean 3 SVD truncations per swap instead
+    of 1, defeating the point of the chain in the first place.
+    """
+    from ..circuits.compiler import QuantumTranspiler
+    from ..circuits.gates import GATE_IDS
+
+    # CCX has no entry in GATE_IDS (it's not a 1- or 2-qubit gate) -- reuse
+    # the same Barenco H/CX/T/Tdg decomposition compiler.py/chunk.py's own
+    # QuantumTranspiler already uses, rather than inventing a second one.
+    # 'swap' is deliberately NOT pre-decomposed here (see id 23 note below).
+    expanded_ops = []
+    for cmd in ops:
+        if str(cmd[0]).lower() == 'ccx':
+            expanded_ops.extend(QuantumTranspiler.decompose_toffoli(*cmd[1:4]))
+        else:
+            expanded_ops.append(cmd)
+
+    rows: List[List[float]] = []
+    for cmd in expanded_ops:
+        name = str(cmd[0]).lower()
+        if name not in GATE_IDS:
+            raise ValueError(
+                f"unknown gate '{cmd[0]}' -- not in GATE_IDS. run_circuit_jit "
+                f"does not silently drop unrecognized gates (same policy as "
+                f"run_circuit_jit_beast_mode, issue #4)."
+            )
+        g_id = float(GATE_IDS[name])
+
+        if g_id < 20:
+            q1 = int(cmd[1])
+            param = float(cmd[2]) if len(cmd) > 2 else 0.0
+            rows.append([g_id, float(q1), float(q1), param, 0.0])
+            continue
+
+        q1 = int(cmd[1])
+        q2 = int(cmd[2])
+        param = float(cmd[3]) if len(cmd) > 3 else 0.0
+
+        if abs(q1 - q2) == 1:
+            lo, hi = (q1, q2) if q1 < q2 else (q2, q1)
+            needs_transpose = q1 > q2
+            rows.append([g_id, float(lo), float(hi), param, 1.0 if needs_transpose else 0.0])
+        else:
+            seq, gate_index, needs_transpose = _expand_nonlocal_2q_positions(q1, q2)
+            for i, (a, b) in enumerate(seq):
+                if i == gate_index:
+                    rows.append([g_id, float(a), float(b), param, 1.0 if needs_transpose else 0.0])
+                else:
+                    rows.append([23.0, float(a), float(b), 0.0, 0.0])  # SWAP
+
+    return rows
+
+
+def _embed_1q_matrix(mat1: np.ndarray, qubit: int, pair: Tuple[int, int]) -> np.ndarray:
+    """Embeds a 2x2 single-qubit matrix into the 4x4 space of a 2-qubit
+    pair (a, b), acting as identity on the other qubit -- kron(mat1, I2)
+    if qubit is the pair's first (more significant) qubit, kron(I2, mat1)
+    otherwise, matching _mps_2q_matrix's own convention (its (4,4)
+    matrices are built with qubit0 as the more significant index before
+    reshaping to (2,2,2,2))."""
+    eye2 = np.eye(2, dtype=mat1.dtype)
+    a, _ = pair
+    return np.kron(mat1, eye2) if qubit == a else np.kron(eye2, mat1)
+
+
+@partial(jax.jit, static_argnums=(2,))
+def _jit_mps_1q_matrix(g_id, param, dtype):
+    """Persistent-compiled-cache wrapper around _mps_1q_matrix, keyed by
+    JAX on (g_id, param, dtype) across the whole process lifetime, not
+    just within one _fuse_compiled_rows call -- without this, each
+    separate run_circuit_jit(fuse_gates=True) call pays a real XLA
+    compile cost again for the same gate/parameter combination, since
+    _mps_1q_matrix's own jax.lax.switch is never wrapped in @jax.jit."""
+    return _mps_1q_matrix(g_id, param, dtype)
+
+
+@partial(jax.jit, static_argnums=(2,))
+def _jit_mps_2q_matrix(g_id, param, dtype):
+    """Same persistent-compiled-cache fix as _jit_mps_1q_matrix, for
+    _mps_2q_matrix."""
+    return _mps_2q_matrix(g_id, param, dtype)
+
+
+def _fuse_compiled_rows(rows: List[List[float]], dtype) -> List[tuple]:
+    """Host-side (pure Python/NumPy, outside any jit region), exact --
+    fuses a chain of consecutive _compile_mps_ops rows into one matrix
+    via matrix multiplication, growing the active qubit pair as needed:
+    a 1-qubit gate on a qubit already inside the current chain's pair
+    gets embedded into that same 4x4 space (kron with identity on the
+    other qubit) rather than breaking the chain -- this is what actually
+    fuses a CX-RZ-CX triple (RZ acts on only one of CX's two qubits, so a
+    naive "exact qubit set match" rule never fuses anything). A chain
+    starting on a single qubit that is later joined by a 2-qubit gate
+    touching it also grows into the 2-qubit space the same way. Any gate
+    touching a qubit OUTSIDE the current chain's pair ends the chain.
+
+    Runs on _compile_mps_ops's OWN output rows, so CCX decomposition and
+    SWAP-chain expansion for non-adjacent gates have already happened --
+    fusion composes correctly with both by construction (a SWAP is just
+    another 2-qubit gate, fusable like any other), verified directly
+    against the eager reference on both a non-adjacent-gate circuit and
+    a CCX circuit, not merely assumed.
+
+    Returns a list of ('1q', q, matrix_2x2) / ('2q', q1, q2, matrix_4x4)
+    entries. transpose_flag is resolved into the matrix itself here, so
+    the fused representation never needs it downstream.
+
+    Memoizes row_matrix by (g_id, param, transpose_flag) -- most rows in
+    a real circuit share the same small set of distinct gate/parameter
+    combinations (e.g. one N=50 TFIM circuit has 985 rows but only 3
+    distinct ones), and each cache miss dispatches a real JAX call
+    (_mps_1q_matrix/_mps_2q_matrix) that is not free to skip. Without
+    this, a real (not tiny-example) circuit pays a full eager JAX
+    dispatch per row on every single run_circuit_jit(fuse_gates=True)
+    call, not just the first -- measured directly: an unmemoized version
+    of this function made "warm" (post-compile) calls nearly as slow as
+    "cold" ones (10.7s vs 12.9s on a 74-gate circuit), because the
+    matrix-reconstruction cost, not recompilation, dominated."""
+    matrix_cache: dict = {}
+
+    def row_matrix(row):
+        g_id, q1, q2, param, transpose_flag = row
+        key = (g_id, param, transpose_flag)
+        cached = matrix_cache.get(key)
+        if cached is not None:
+            return cached, (int(q1), int(q2)) if g_id >= 20 else (int(q1),)
+        g_id_arr = jnp.asarray(g_id).astype(jnp.int32)
+        if g_id >= 20:
+            mat4 = np.asarray(_jit_mps_2q_matrix(g_id_arr, jnp.asarray(param), dtype))
+            if transpose_flag > 0.5:
+                mat4 = np.transpose(mat4, (1, 0, 3, 2))
+            mat = mat4.reshape(4, 4)
+        else:
+            mat = np.asarray(_jit_mps_1q_matrix(g_id_arr, jnp.asarray(param), dtype))
+        matrix_cache[key] = mat
+        return mat, (int(q1), int(q2)) if g_id >= 20 else (int(q1),)
+
+    fused = []
+    i = 0
+    n = len(rows)
+    while i < n:
+        mat, qubits = row_matrix(rows[i])
+        pair = qubits if len(qubits) == 2 else None
+        active_qubit = qubits[0]
+        j = i + 1
+        while j < n:
+            nmat, nqubits = row_matrix(rows[j])
+            n_is_2q = len(nqubits) == 2
+            if pair is None:
+                if not n_is_2q:
+                    if nqubits[0] != active_qubit:
+                        break
+                    mat = nmat @ mat
+                    j += 1
+                    continue
+                if active_qubit not in nqubits:
+                    break
+                pair = nqubits
+                mat = _embed_1q_matrix(mat, active_qubit, pair)
+                mat = nmat @ mat
+                j += 1
+                continue
+            if n_is_2q:
+                if nqubits != pair:
+                    break
+            else:
+                if nqubits[0] not in pair:
+                    break
+                nmat = _embed_1q_matrix(nmat, nqubits[0], pair)
+            mat = nmat @ mat
+            j += 1
+        if pair is not None:
+            fused.append(('2q', pair[0], pair[1], mat.reshape(2, 2, 2, 2)))
+        else:
+            fused.append(('1q', active_qubit, mat))
+        i = j
+    return fused
+
+
+def _fused_entries_to_arrays(fused: List[tuple], dtype):
+    """Turns _fuse_compiled_rows's output into the parallel arrays
+    jax.lax.scan needs -- every step carries BOTH a 1q and a 2q matrix
+    slot (one is an unused identity placeholder) so shapes stay uniform
+    across the whole scanned sequence, same principle _mps_1q_matrix/
+    _mps_2q_matrix's own switch already relies on (all branches traced,
+    only one path's real work executed)."""
+    is_2q, q1s, q2s, mats2q, mats1q = [], [], [], [], []
+    eye2 = np.eye(2, dtype=dtype)
+    eye4 = np.eye(4, dtype=dtype).reshape(2, 2, 2, 2)
+    for entry in fused:
+        if entry[0] == '2q':
+            _, a, b, mat = entry
+            is_2q.append(True)
+            q1s.append(a)
+            q2s.append(b)
+            mats2q.append(mat)
+            mats1q.append(eye2)
+        else:
+            _, a, mat = entry
+            is_2q.append(False)
+            q1s.append(a)
+            q2s.append(0)
+            mats2q.append(eye4)
+            mats1q.append(mat)
+    return (jnp.asarray(is_2q), jnp.asarray(q1s, dtype=jnp.int32), jnp.asarray(q2s, dtype=jnp.int32),
+            jnp.asarray(np.stack(mats2q)), jnp.asarray(np.stack(mats1q)))
+
+
+_SVD_BUCKETS = (2, 4, 8, 16, 32, 64, 128, 256, 512, 1024)
+
+
+def _bucket_sizes(max_bond: int) -> Tuple[int, ...]:
+    """Smallest-to-largest candidate SVD sizes for one MPSSimulator
+    instance's whole lifetime -- always ending at max_bond itself (added
+    if not already a power of 2 in _SVD_BUCKETS), so every real bond
+    dimension up to the hard cap has some bucket that fits it."""
+    buckets = tuple(b for b in _SVD_BUCKETS if b <= max_bond)
+    if not buckets or buckets[-1] != max_bond:
+        buckets = buckets + (max_bond,)
+    return buckets
+
+
+def _build_mps_runner(n_qubits: int, max_bond: int, eps: float, jsd_budget: float):
+    """Factory: builds and returns a single @jax.jit-compiled closure that
+    runs an entire pre-compiled MPS circuit via jax.lax.scan -- same
+    factory pattern as chunk.py's _build_multi_chunk_runner, closing over
+    the static n_qubits/max_bond/eps/jsd_budget (fixed for one
+    MPSSimulator instance's whole lifetime, so the returned closure is
+    built once and cached on the instance, never rebuilt per call).
+
+    step's branch_1q returns a placeholder diag tuple (zeros) for 1-qubit
+    steps -- run_circuit_jit filters these out by g_id (>=20) before using
+    the per-step history, same as _bond_history/jsd_per_bond only ever
+    growing on 2-qubit gates in the eager path.
+
+    2-qubit steps no longer always run the SVD at a fixed max_bond*2
+    size -- previously theta_mat.reshape(max_bond*2, 2*max_bond) was
+    ALWAYS square and ALWAYS a full economy SVD at that size, regardless
+    of the real/masked bond dimension (confirmed by direct measurement:
+    raising max_bond 64->128 measurably slowed every gate down even when
+    the real bond stayed ~4-8 on the same circuit). Instead, a small fixed
+    set of candidate sizes (_bucket_sizes) is dispatched via
+    jax.lax.switch -- same JIT-compatible mechanism _mps_1q_matrix/
+    _mps_2q_matrix already use for gate-ID dispatch, applied here to SVD
+    size instead. Verified correct across ~92 single-gate configurations
+    plus full multi-gate circuit fidelity (0.999999999998+) against this
+    module's own eager path, and 68.80x-73.96x faster on CPU / 2.74x on
+    GPU on a real N=50 TFIM Trotter circuit -- see Dense-Evolution-
+    Discovery's mps_bucketed_svd_optimization experiment for the full
+    validation, including a real bug found and fixed there before
+    promotion (see the bound formula's own comment below).
+
+    Bucket selection needs a provable (not heuristic) bound on how big B
+    must be, computed BEFORE the SVD from already-known real bond sizes
+    (real_chi, threaded through the scan carry alongside gammas/lambdas):
+    a 2-qubit gate acting at a cut can increase that cut's Schmidt rank by
+    at most a factor of d^2=4 (d=2, local physical dimension) -- theta_mat's
+    own matrix rank is bounded by min(chi_l_real*2, chi_r_real*2), a
+    mathematical fact about rank <= min(rows, cols), independent of the
+    PRE-gate middle bond. That output-side bound alone is not sufficient,
+    though: g1/g2/lam_m get sliced to size B BEFORE the einsum contracts
+    the middle bond away, so B must ALSO be at least as large as the
+    current middle bond (chi_m_real) and both outer bonds themselves, or
+    real (nonzero) Schmidt weight already on those bonds gets silently
+    dropped from the contraction -- not merely under-grown. Both
+    requirements combined: max(chi_l_real, chi_r_real, chi_m_real,
+    min(chi_l_real*2, chi_r_real*2))."""
+
+    buckets = _bucket_sizes(max_bond)
+    bucket_arr = jnp.array(buckets)
+
+    def step(carry, row):
+        gammas, lambdas, real_chi = carry
+        dtype = gammas.dtype
+
+        g_id = row[0].astype(jnp.int32)
+        q1 = row[1].astype(jnp.int32)
+        q2 = row[2].astype(jnp.int32)
+        param = row[3]
+        transpose_flag = row[4] > 0.5
+
+        def branch_1q(c):
+            gammas_, lambdas_, real_chi_ = c
+            gate_1q = _mps_1q_matrix(g_id, param, dtype)
+            new_g = jnp.einsum('ij,ljr->lir', gate_1q, gammas_[q1])
+            new_carry = (gammas_.at[q1].set(new_g), lambdas_, real_chi_)
+            real_dtype = _real_dtype_for(dtype)
+            diag = (jnp.asarray(0, dtype=jnp.int32), jnp.asarray(0.0, dtype=real_dtype),
+                    jnp.asarray(0.0, dtype=real_dtype), jnp.asarray(0.0, dtype=real_dtype))
+            return new_carry, diag
+
+        def branch_2q(c):
+            # Traced unconditionally only when this branch is taken --
+            # jax.lax.cond (not jnp.where) skips the untaken branch's
+            # compute entirely, same choice compiler.py's do_2q makes for
+            # its own outer 1q/2q split: the SVD here is real work, worth
+            # skipping for the common case of a 1-qubit gate.
+            gammas_, lambdas_, real_chi_ = c
+            gate_2q = _mps_2q_matrix(g_id, param, dtype)
+            gate_2q = jnp.where(transpose_flag, jnp.transpose(gate_2q, (1, 0, 3, 2)), gate_2q)
+
+            chi_l_real = real_chi_[q1]
+            chi_m_real = real_chi_[q2]      # current middle bond, BEFORE this gate
+            chi_r_real = real_chi_[q2 + 1]
+            input_min = jnp.maximum(jnp.maximum(chi_l_real, chi_r_real), chi_m_real)
+            output_bound = jnp.minimum(chi_l_real * 2, chi_r_real * 2)
+            bound = jnp.maximum(input_min, output_bound)
+            ge_mask = bucket_arr >= bound
+            bucket_idx = jnp.where(jnp.any(ge_mask), jnp.argmax(ge_mask), len(buckets) - 1)
+
+            def make_branch(B):
+                def branch_fn(_):
+                    g1 = gammas_[q1][:B, :, :B]
+                    g2 = gammas_[q2][:B, :, :B]
+                    lam_l = lambdas_[q1][:B]
+                    lam_m = lambdas_[q2][:B]
+                    lam_r = lambdas_[q2 + 1][:B]
+                    # Full Vidal update (prog.txt P0 fix): both outer
+                    # Lambdas attached, same convention as the eager
+                    # apply_gate_2q path -- see that method's docstring.
+                    theta = jnp.einsum('l,lik,k,kjr,r->lijr', lam_l, g1, lam_m, g2, lam_r)
+                    theta_new = jnp.einsum('abcd,ecdf->eabf', gate_2q, theta)
+                    theta_mat = theta_new.reshape(B * 2, 2 * B)
+
+                    U, S, Vh = jnp.linalg.svd(theta_mat, full_matrices=False)
+                    chi_new, jsd_val = _vectorized_chi_search_jax(S, eps, jsd_budget, min(B, max_bond))
+                    col_mask = jnp.arange(B) < chi_new
+
+                    norm_full = jnp.sqrt(jnp.sum(S ** 2) + 1e-30)
+                    S_norm_full = S / (norm_full + 1e-30)
+                    trunc_err = jnp.sqrt(jnp.sum(jnp.where(jnp.arange(2 * B) >= chi_new, S_norm_full ** 2, 0.0)))
+
+                    S_kept_masked = jnp.where(col_mask, S[:B], 0.0)
+                    kept_norm = jnp.sqrt(jnp.sum(S_kept_masked ** 2) + 1e-30)
+                    S_fixed = jnp.where(col_mask, S_kept_masked / (kept_norm + 1e-30), 0.0)
+
+                    lam_l_inv = jnp.where(lam_l > eps, 1.0 / lam_l, 0.0)
+                    lam_r_inv = jnp.where(lam_r > eps, 1.0 / lam_r, 0.0)
+
+                    U_masked = jnp.where(col_mask[None, :], U[:, :B], 0.0)
+                    Vh_masked = jnp.where(col_mask[:, None], Vh[:B, :], 0.0)
+
+                    new_g1 = jnp.einsum('l,lir->lir', lam_l_inv, U_masked.reshape(B, 2, B))
+                    new_g2 = jnp.einsum('ljr,r->ljr', Vh_masked.reshape(B, 2, B), lam_r_inv)
+
+                    p_dist = S_fixed ** 2  # already unit-norm (see eager _svd_truncate)
+                    ee = -jnp.sum(jnp.where(p_dist > 1e-20, p_dist * jnp.log2(jnp.where(p_dist > 1e-20, p_dist, 1.0)), 0.0))
+
+                    real_dtype = _real_dtype_for(dtype)
+                    return (_pad_gamma(new_g1, max_bond), _pad_gamma(new_g2, max_bond),
+                            _pad_lambda(S_fixed, max_bond), chi_new.astype(jnp.int32),
+                            jsd_val.astype(real_dtype), trunc_err.astype(real_dtype), ee.astype(real_dtype))
+
+                return branch_fn
+
+            branches = [make_branch(B) for B in buckets]
+            new_g1_p, new_g2_p, S_fixed_p, chi_new, jsd_val, trunc_err, ee = jax.lax.switch(
+                bucket_idx, branches, operand=None)
+
+            new_gammas = gammas_.at[q1].set(new_g1_p).at[q2].set(new_g2_p)
+            new_lambdas = lambdas_.at[q2].set(S_fixed_p)
+            new_real_chi = real_chi_.at[q2].set(chi_new)
+
+            new_carry = (new_gammas, new_lambdas, new_real_chi)
+            diag = (chi_new, jsd_val, trunc_err, ee)
+            return new_carry, diag
+
+        is_2q = g_id >= 20
+        new_carry, diag = jax.lax.cond(is_2q, branch_2q, branch_1q, carry)
+        return new_carry, diag
+
+    @jax.jit
+    def run(gammas, lambdas, real_chi, compiled_ops):
+        (final_gammas, final_lambdas, final_real_chi), diag = jax.lax.scan(
+            step, (gammas, lambdas, real_chi), compiled_ops)
+        return final_gammas, final_lambdas, final_real_chi, diag
+
+    return run
+
+
+def _build_fused_mps_runner(n_qubits: int, max_bond: int, eps: float, jsd_budget: float):
+    """Factory for run_circuit_jit(ops, fuse_gates=True) -- same bucketed-
+    SVD dispatch as _build_mps_runner, except every step's gate matrix is
+    taken directly from a pre-fused matrix stream (built host-side by
+    _fuse_compiled_rows/_fused_entries_to_arrays from _compile_mps_ops's
+    own output) instead of being reconstructed inside the JIT via
+    _mps_1q_matrix/_mps_2q_matrix -- gate-ID dispatch is removed from the
+    traced program entirely, not just extended. Fusing consecutive
+    same-qubit-pair gates into one matrix means fewer, larger scan steps:
+    a real, measured ~2x GPU speedup on top of the bucketed dispatch
+    alone (see Dense-Evolution-Discovery's mps_gate_blocking_redesign_v2
+    experiment), on top of jax.lax.scan's own per-step GPU dispatch
+    overhead being the real bottleneck the bucketed dispatch alone
+    couldn't remove.
+
+    Returns the same (chi_new, jsd_val, trunc_err, entanglement_entropy)
+    diagnostic tuple _build_mps_runner does, one entry per FUSED step
+    rather than per original gate -- run_circuit_jit populates
+    self._bond_history/jsd_per_bond/truncation_errors/entanglement_entropy
+    from this at that coarser granularity when fuse_gates=True, an
+    explicit, documented trade-off (see MPSSimulator.run_circuit_jit's
+    own docstring), not a silent behavior change."""
+    buckets = _bucket_sizes(max_bond)
+    bucket_arr = jnp.array(buckets)
+
+    def step(carry, xs):
+        gammas, lambdas, real_chi = carry
+        dtype = gammas.dtype
+        is_2q, q1, q2, mat2q, mat1q = xs
+        q1 = q1.astype(jnp.int32)
+        q2 = q2.astype(jnp.int32)
+        real_dtype = _real_dtype_for(dtype)
+
+        def branch_1q(c):
+            gammas_, lambdas_, real_chi_ = c
+            new_g = jnp.einsum('ij,ljr->lir', mat1q, gammas_[q1])
+            new_carry = (gammas_.at[q1].set(new_g), lambdas_, real_chi_)
+            diag = (jnp.asarray(0, dtype=jnp.int32), jnp.asarray(0.0, dtype=real_dtype),
+                    jnp.asarray(0.0, dtype=real_dtype), jnp.asarray(0.0, dtype=real_dtype))
+            return new_carry, diag
+
+        def branch_2q(c):
+            gammas_, lambdas_, real_chi_ = c
+            gate_2q = mat2q
+
+            chi_l_real = real_chi_[q1]
+            chi_m_real = real_chi_[q2]
+            chi_r_real = real_chi_[q2 + 1]
+            input_min = jnp.maximum(jnp.maximum(chi_l_real, chi_r_real), chi_m_real)
+            output_bound = jnp.minimum(chi_l_real * 2, chi_r_real * 2)
+            bound = jnp.maximum(input_min, output_bound)
+            ge_mask = bucket_arr >= bound
+            bucket_idx = jnp.where(jnp.any(ge_mask), jnp.argmax(ge_mask), len(buckets) - 1)
+
+            def make_branch(B):
+                def branch_fn(_):
+                    g1 = gammas_[q1][:B, :, :B]
+                    g2 = gammas_[q2][:B, :, :B]
+                    lam_l = lambdas_[q1][:B]
+                    lam_m = lambdas_[q2][:B]
+                    lam_r = lambdas_[q2 + 1][:B]
+                    theta = jnp.einsum('l,lik,k,kjr,r->lijr', lam_l, g1, lam_m, g2, lam_r)
+                    theta_new = jnp.einsum('abcd,ecdf->eabf', gate_2q, theta)
+                    theta_mat = theta_new.reshape(B * 2, 2 * B)
+
+                    U, S, Vh = jnp.linalg.svd(theta_mat, full_matrices=False)
+                    chi_new, jsd_val = _vectorized_chi_search_jax(S, eps, jsd_budget, min(B, max_bond))
+                    col_mask = jnp.arange(B) < chi_new
+
+                    norm_full = jnp.sqrt(jnp.sum(S ** 2) + 1e-30)
+                    S_norm_full = S / (norm_full + 1e-30)
+                    trunc_err = jnp.sqrt(jnp.sum(jnp.where(jnp.arange(2 * B) >= chi_new, S_norm_full ** 2, 0.0)))
+
+                    S_kept_masked = jnp.where(col_mask, S[:B], 0.0)
+                    kept_norm = jnp.sqrt(jnp.sum(S_kept_masked ** 2) + 1e-30)
+                    S_fixed = jnp.where(col_mask, S_kept_masked / (kept_norm + 1e-30), 0.0)
+
+                    lam_l_inv = jnp.where(lam_l > eps, 1.0 / lam_l, 0.0)
+                    lam_r_inv = jnp.where(lam_r > eps, 1.0 / lam_r, 0.0)
+
+                    U_masked = jnp.where(col_mask[None, :], U[:, :B], 0.0)
+                    Vh_masked = jnp.where(col_mask[:, None], Vh[:B, :], 0.0)
+
+                    new_g1 = jnp.einsum('l,lir->lir', lam_l_inv, U_masked.reshape(B, 2, B))
+                    new_g2 = jnp.einsum('ljr,r->ljr', Vh_masked.reshape(B, 2, B), lam_r_inv)
+
+                    p_dist = S_fixed ** 2
+                    ee = -jnp.sum(jnp.where(p_dist > 1e-20, p_dist * jnp.log2(jnp.where(p_dist > 1e-20, p_dist, 1.0)), 0.0))
+
+                    return (_pad_gamma(new_g1, max_bond), _pad_gamma(new_g2, max_bond),
+                            _pad_lambda(S_fixed, max_bond), chi_new.astype(jnp.int32),
+                            jsd_val.astype(real_dtype), trunc_err.astype(real_dtype), ee.astype(real_dtype))
+
+                return branch_fn
+
+            branches = [make_branch(B) for B in buckets]
+            new_g1_p, new_g2_p, S_fixed_p, chi_new, jsd_val, trunc_err, ee = jax.lax.switch(
+                bucket_idx, branches, operand=None)
+
+            new_gammas = gammas_.at[q1].set(new_g1_p).at[q2].set(new_g2_p)
+            new_lambdas = lambdas_.at[q2].set(S_fixed_p)
+            new_real_chi = real_chi_.at[q2].set(chi_new)
+
+            new_carry = (new_gammas, new_lambdas, new_real_chi)
+            diag = (chi_new, jsd_val, trunc_err, ee)
+            return new_carry, diag
+
+        new_carry, diag = jax.lax.cond(is_2q, branch_2q, branch_1q, carry)
+        return new_carry, diag
+
+    @jax.jit
+    def run(gammas, lambdas, real_chi, xs):
+        (final_gammas, final_lambdas, final_real_chi), diag = jax.lax.scan(
+            step, (gammas, lambdas, real_chi), xs)
+        return final_gammas, final_lambdas, final_real_chi, diag
+
+    return run
+
+
+class MPSSimulator:
+    """
+    Matrix Product State simulator with adaptive SVD-truncated bond
+    dimension (JSD-budget driven), no lossy post-truncation quantization.
+    JAX-backed core (einsum, SVD).
+
+    Parameters
+    ----------
+    n_qubits    : int
+    max_bond    : int   -- hard cap on bond dimension chi
+    svd_cutoff  : float or None -- singular values below this are dropped
+                          outright. None (default) resolves to a value
+                          appropriate for the dtype this instance actually
+                          runs at: ~1e-12 for complex128, ~1e-6 (roughly
+                          10x float32's own machine epsilon) for complex64
+                          -- a fixed 1e-12 in complex64 sits below that
+                          dtype's noise floor, so numerical noise gets
+                          counted as real Schmidt weight and chi never
+                          shrinks below max_bond regardless of jsd_budget.
+                          An explicitly passed value always wins verbatim,
+                          never rescaled.
+    jsd_budget  : float -- max tolerated Jensen-Shannon distance between the
+                          full and truncated singular-value distributions
+                          at each cut; chi is grown by 1 until satisfied
+                          or max_bond is hit.
+    use_float32 : bool or None -- None (default) follows the process-wide
+                          jax_enable_x64 flag, same convention as this
+                          module always used (see module docstring).
+                          True forces complex64 (and the complex64-
+                          appropriate svd_cutoff default) even if x64 is
+                          enabled. False requests complex128 -- since
+                          complex128 arrays don't exist in JAX at all
+                          unless the process-wide flag is on, this calls
+                          the same lazy ensure_x64() DenseSVSimulator
+                          uses, mirroring its own use_float32=False
+                          handling. That call is a no-op if precision was
+                          already pinned via set_precision(), same
+                          deference DenseSVSimulator gives it; the flag is
+                          re-read afterward rather than assumed, so
+                          dtype/eps stay consistent with whatever
+                          precision is really active even in that
+                          pinned-False edge case.
+    """
+
+    def __init__(
+        self,
+        n_qubits: int,
+        max_bond: int = 64,
+        svd_cutoff: Optional[float] = None,
+        jsd_budget: float = 1e-5,
+        use_float32: Optional[bool] = None,
+    ):
+        self.n = n_qubits
+        self.chi = max_bond
+        if use_float32 is None:
+            x64_active = jax.config.jax_enable_x64
+        elif use_float32:
+            x64_active = False
+        else:
+            ensure_x64()
+            x64_active = jax.config.jax_enable_x64
+        dtype = jnp.complex128 if x64_active else jnp.complex64
+        self.eps = svd_cutoff if svd_cutoff is not None else (1e-12 if x64_active else 1e-6)
+        self.jsd_budget = jsd_budget
+
+        self.gammas: List[jnp.ndarray] = []
+        self.lambdas: List[jnp.ndarray] = [jnp.ones(1)] * (n_qubits + 1)
+        # Real (non-max_bond-padded) bond dimension at every cut, kept in
+        # sync by BOTH the eager path (apply_gate_2q, below) and
+        # run_circuit_jit -- needed by the bucketed-SVD dispatch in
+        # _build_mps_runner to pick a provably-sufficient bucket size
+        # without ever inferring it from zero-counting (see that
+        # function's own docstring for why that would be unreliable).
+        self._real_chi: np.ndarray = np.ones(n_qubits + 1, dtype=np.int64)
+
+        self.truncation_errors: List[float] = []
+        self.jsd_per_bond: List[float] = []
+        self.entanglement_entropy = np.zeros(max(n_qubits - 1, 0))
+        self._bond_history: List[int] = []
+        # Counts truncations where max_bond was hit before jsd_budget could
+        # be satisfied -- the while loop below exits silently in that case,
+        # and avg_JSD (a mean over all steps) can look deceptively low even
+        # when the final contracted state is badly wrong (verified: TVD
+        # ~0.97 against DenseSVSimulator on an 8-qubit/15-layer entangling
+        # circuit with max_bond=2, while avg_JSD read 0.0534).
+        self.budget_violations: int = 0
+
+        # Cached compiled closure for run_circuit_jit -- built lazily on
+        # first use (self.n/self.chi/self.eps/self.jsd_budget are fixed
+        # for this instance's lifetime), never rebuilt per call. Same
+        # caching pattern as Chunk.__init__'s self._multi_chunk_runner.
+        self._mps_runner = None
+        self._fused_mps_runner = None
+
+        for _ in range(n_qubits):
+            g = jnp.zeros((1, 2, 1), dtype=dtype)
+            g = g.at[0, 0, 0].set(1.0)
+            self.gammas.append(g)
+
+    # gate 1q
+    def apply_gate_1q(self, gate: jnp.ndarray, qubit: int) -> None:
+        """O(chi^2) -- updates only Gamma[qubit]."""
+        gate = jnp.asarray(gate)
+        self.gammas[qubit] = jnp.einsum("ij,ljr->lir", gate, self.gammas[qubit])
+
+    # core: plain adaptive SVD truncation (no quantization)
+    def _svd_truncate(
+        self, theta_mat: jnp.ndarray
+    ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, float, float]:
+        """SVD + adaptive chi search, plus the two corrections from prog.txt
+        P0: trunc_err/entanglement_entropy need the true (normalized)
+        Schmidt spectrum, and the kept singular values are renormalized to
+        unit norm after truncation (same convention Vidal's iTEBD update
+        uses -- the state's overall norm is the product of every bond's
+        Lambda norm, each already implicitly 1 by induction, so explicitly
+        restoring that here keeps the whole chain normalized without ever
+        relying on a global re-norm in contract_to_statevector)."""
+        U, S, Vh = jnp.linalg.svd(theta_mat, full_matrices=False)
+
+        # Was a Python while loop incrementing chi_new one at a time,
+        # host-syncing (float()/int() casts) every iteration -- the
+        # eager path never used the vectorized search already written
+        # for the JIT path (_vectorized_chi_search / ..._jax), despite
+        # that search being verified to replicate the while loop's
+        # chi_new/jsd_val exactly (0 mismatches across 171 real calls,
+        # including this budget-violation fallback -- see that
+        # function's own docstring).
+        chi_new, jsd_val = _vectorized_chi_search(S, self.eps, self.jsd_budget, self.chi)
+        max_possible = min(len(S), self.chi)
+
+        if jsd_val > self.jsd_budget:
+            # Loop exited because chi_new hit max_possible (== max_bond, in
+            # the common case where the bond isn't already capped by the
+            # SVD's own rank), not because jsd_budget was satisfied.
+            if self.budget_violations == 0:
+                warnings.warn(
+                    f"MPSSimulator: bond dimension capped at max_bond={self.chi}, "
+                    f"jsd_budget={self.jsd_budget:.1e} not honored "
+                    f"(jsd={jsd_val:.2e}) -- results may be unreliable, "
+                    f"consider raising max_bond.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+            self.budget_violations += 1
+
+        norm_full = float(jnp.sqrt(jnp.sum(S ** 2) + 1e-30))
+        S_norm_full = S / (norm_full + 1e-30)
+        trunc_err = float(jnp.sqrt(jnp.sum(S_norm_full[chi_new:] ** 2))) if len(S) > chi_new else 0.0
+        self.truncation_errors.append(trunc_err)
+
+        S_kept = S[:chi_new]
+        kept_norm = float(jnp.sqrt(jnp.sum(S_kept ** 2) + 1e-30))
+        S_kept_renorm = S_kept / (kept_norm + 1e-30)
+
+        return U[:, :chi_new], S_kept_renorm, Vh[:chi_new, :], trunc_err, jsd_val
+
+    # gate 2q adjacent
+    def apply_gate_2q(self, gate_2q: jnp.ndarray, q1: int, q2: int) -> None:
+        """2-qubit gate with adaptive SVD truncation. O(chi^3).
+
+        Vidal's full two-site update (prog.txt P0 fix): theta is built from
+        BOTH outer Lambdas (Lambda[q1], Lambda[q2+1]) as well as the middle
+        one, not just the middle one -- so its singular values are the true
+        global Schmidt coefficients at this cut, not an artifact of the
+        local 2-site reduced state. New Gamma tensors are recovered by
+        dividing the outer Lambdas back out (regularized: entries at or
+        below svd_cutoff map to a zero inverse instead of blowing up --
+        exactly the padded/zero entries in the JIT path's fixed-size
+        arrays, and the trivial size-1 boundary Lambda everywhere else)."""
+        gate_2q = jnp.asarray(gate_2q)
+        if abs(q1 - q2) != 1:
+            self._apply_nonlocal_2q(gate_2q, q1, q2)
+            return
+        if q1 > q2:
+            q1, q2 = q2, q1
+            gate_2q = jnp.transpose(gate_2q, (1, 0, 3, 2))
+
+        lam_l = self.lambdas[q1]
+        g1 = self.gammas[q1]
+        lam_m = self.lambdas[q2]
+        g2 = self.gammas[q2]
+        lam_r = self.lambdas[q2 + 1]
+
+        theta = jnp.einsum("l,lik,k,kjr,r->lijr", lam_l, g1, lam_m, g2, lam_r)
+        chiL, d1, d2, chiR = theta.shape
+
+        theta_new = jnp.einsum("abcd,ecdf->eabf", gate_2q, theta)
+        theta_mat = theta_new.reshape(chiL * d1, d2 * chiR)
+
+        U_t, S_t, Vh_t, trunc_err, jsd_val = self._svd_truncate(theta_mat)
+        chi_new = len(S_t)
+
+        lam_l_inv = jnp.where(lam_l > self.eps, 1.0 / lam_l, 0.0)
+        lam_r_inv = jnp.where(lam_r > self.eps, 1.0 / lam_r, 0.0)
+
+        new_g1 = jnp.einsum("l,lir->lir", lam_l_inv, U_t.reshape(chiL, d1, chi_new))
+        new_g2 = jnp.einsum("ljr,r->ljr", Vh_t.reshape(chi_new, d2, chiR), lam_r_inv)
+
+        p_dist = S_t ** 2  # already unit-norm (see _svd_truncate)
+        mask = p_dist > 1e-20
+        ee = float(-jnp.sum(jnp.where(mask, p_dist * jnp.log2(jnp.where(mask, p_dist, 1.0)), 0.0)))
+        if q1 < len(self.entanglement_entropy):
+            self.entanglement_entropy[q1] = ee
+
+        self.lambdas[q2] = S_t
+        self.gammas[q1] = new_g1
+        self.gammas[q2] = new_g2
+        self._real_chi[q2] = chi_new
+
+        self._bond_history.append(chi_new)
+        self.jsd_per_bond.append(jsd_val)
+
+    def _apply_nonlocal_2q(self, gate_2q: jnp.ndarray, q1: int, q2: int) -> None:
+        """Non-adjacent 2-qubit gate via a SWAP chain down to adjacent.
+
+        The SWAP chain always ends up applying the gate at (target, target+1)
+        with target = min(q1, q2) first -- so without normalizing here, a
+        caller passing q1 > q2 (e.g. apply_cx(ctrl=3, tgt=1)) would silently
+        have its control/target roles swapped for asymmetric gates like CNOT.
+        Same fix as apply_gate_2q's adjacent-qubit branch just above, applied
+        before the swap chain runs so it always sees q1 < q2."""
+        if q1 > q2:
+            q1, q2 = q2, q1
+            gate_2q = jnp.transpose(gate_2q, (1, 0, 3, 2))
+        swap = jnp.array(
+            [[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=complex
+        ).reshape(2, 2, 2, 2)
+        target = min(q1, q2)
+        for q in range(max(q1, q2) - 1, target, -1):
+            self.apply_gate_2q(swap, q, q + 1)
+        self.apply_gate_2q(gate_2q, target, target + 1)
+        for q in range(target + 1, max(q1, q2)):
+            self.apply_gate_2q(swap, q, q + 1)
+
+    # gate shortcuts
+    def apply_cx(self, ctrl: int, tgt: int) -> None:
+        cx = jnp.array(
+            [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]], dtype=complex
+        ).reshape(2, 2, 2, 2)
+        self.apply_gate_2q(cx, ctrl, tgt)
+
+    def apply_cz(self, ctrl: int, tgt: int) -> None:
+        cz = jnp.diag(jnp.array([1, 1, 1, -1])).astype(complex).reshape(2, 2, 2, 2)
+        self.apply_gate_2q(cz, ctrl, tgt)
+
+    def apply_swap(self, q1: int, q2: int) -> None:
+        sw = jnp.array(
+            [[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]], dtype=complex
+        ).reshape(2, 2, 2, 2)
+        self.apply_gate_2q(sw, q1, q2)
+
+    def apply_ccx(self, c1: int, c2: int, tgt: int) -> None:
+        """Toffoli via standard T-gate decomposition (all 1q/2q gates)."""
+        inv2 = 1.0 / np.sqrt(2.0)
+        h = inv2 * jnp.array([[1, 1], [1, -1]], dtype=complex)
+        t = jnp.array([[1, 0], [0, jnp.exp(1j * jnp.pi / 4)]], dtype=complex)
+        tdg = jnp.array([[1, 0], [0, jnp.exp(-1j * jnp.pi / 4)]], dtype=complex)
+
+        self.apply_gate_1q(h, tgt)
+        self.apply_cx(c2, tgt)
+        self.apply_gate_1q(tdg, tgt)
+        self.apply_cx(c1, tgt)
+        self.apply_gate_1q(t, tgt)
+        self.apply_cx(c2, tgt)
+        self.apply_gate_1q(tdg, tgt)
+        self.apply_cx(c1, tgt)
+        self.apply_gate_1q(t, c2)
+        self.apply_gate_1q(t, tgt)
+        self.apply_gate_1q(h, tgt)
+        self.apply_cx(c1, c2)
+        self.apply_gate_1q(t, c1)
+        self.apply_gate_1q(tdg, c2)
+        self.apply_cx(c1, c2)
+
+    # contraction to a full statevector -- O(2**n), n <= ~24 only
+    def contract_to_statevector(self) -> jnp.ndarray:
+        if self.n > 24:
+            raise MemoryError(
+                f"contract_to_statevector at {self.n} qubits would need "
+                f"~{2**self.n * 16 / 1e9:.1f} GB -- use "
+                f"get_probabilities_sampled or get_top_k_probable_states "
+                f"instead for n > 24."
+            )
+        # Index [0] rather than .squeeze(axis=0)/.squeeze(axis=-1): the
+        # boundary axis has size 1 in the eager (unpadded) representation
+        # (squeeze and [0]-indexing are identical there), but size
+        # max_bond after run_circuit_jit's fixed-shape padding -- the real
+        # boundary content is always at index 0 by construction either
+        # way (see _pad_gamma), so [0]-indexing is the strict
+        # generalization that works for both, not a behavior change for
+        # the pre-existing eager path.
+        result = self.gammas[0][0, :, :]
+        for i in range(1, self.n):
+            lam = self.lambdas[i]
+            g = self.gammas[i]
+            lg = jnp.einsum("k,kir->kir", lam, g)
+            result = jnp.tensordot(result, lg, axes=([-1], [0]))
+            result = result.reshape(-1, result.shape[-1])
+        sv = result[..., 0]
+        norm = jnp.linalg.norm(sv)
+        return sv / (norm + 1e-15)
+
+    # sequential sampling -- O(chi^2) per qubit per sample, never
+    # materializes a (2**n,) array. The only way to get results for
+    # n_qubits beyond ~24. Sampling decisions themselves stay on host
+    # (np.random.Generator): each step needs a concrete probability to
+    # branch on, not a traced value, so this loop isn't a jax.jit
+    # candidate as a whole regardless of backend.
+    def _sample_bitstring(self, rng: np.random.Generator) -> List[int]:
+        bits = []
+        state = jnp.ones(1, dtype=self.gammas[0].dtype)
+        for i in range(self.n):
+            g = self.gammas[i]
+            lam = self.lambdas[i + 1] if i < self.n - 1 else jnp.ones(g.shape[2])
+            p0v = jnp.einsum("l,lr->r", state, g[:, 0, :])
+            p1v = jnp.einsum("l,lr->r", state, g[:, 1, :])
+            p0l = p0v * lam
+            p1l = p1v * lam
+            p0 = float(jnp.real(jnp.dot(p0l, jnp.conj(p0l))))
+            p1 = float(jnp.real(jnp.dot(p1l, jnp.conj(p1l))))
+            norm = p0 + p1 + 1e-15
+            bit = 0 if rng.random() < p0 / norm else 1
+            bits.append(bit)
+            state = p0l if bit == 0 else p1l
+            state = state / (jnp.linalg.norm(state) + 1e-15)
+        return bits
+
+    def get_probabilities_sampled(
+        self, n_samples: int = 100_000, seed: Optional[int] = None
+    ) -> dict:
+        """Returns a {bitstring: empirical_probability} dict from n_samples
+        sequential draws -- the only entry point safe for n_qubits > 24."""
+        from collections import Counter
+
+        rng = np.random.default_rng(seed)
+        counts: Counter = Counter()
+        for _ in range(n_samples):
+            bits = self._sample_bitstring(rng)
+            counts["".join(map(str, bits))] += 1
+        return {bitstr: c / n_samples for bitstr, c in counts.items()}
+
+    # ──────────────────────────────────────────────
+    # approximate top-k extraction -- see module docstring point 2.
+    # ──────────────────────────────────────────────
+    def get_top_k_probable_states(self, k: int = 128) -> Tuple[np.ndarray, np.ndarray]:
+        """Greedy beam search (beam width k) for approximately-most-probable
+        basis states, without ever contracting to a full statevector.
+
+        Returns (indices, probabilities): indices are computational-basis
+        integers, probabilities are exact for the states found (not
+        approximated), sorted descending. Recall of the TRUE top states
+        improves with k but is not guaranteed for any fixed k -- see the
+        module docstring."""
+        paths: List[Tuple[int, jnp.ndarray]] = [(0, jnp.array([1.0 + 0.0j]))]
+        for i in range(self.n):
+            candidates = []
+            gamma = self.gammas[i]
+            lam = self.lambdas[i + 1] if (i + 1) < len(self.lambdas) else jnp.ones(gamma.shape[2])
+            for idx_p, vec_p in paths:
+                for bit in (0, 1):
+                    new_vec = jnp.einsum("l,lr->r", vec_p, gamma[:, bit, :]) * lam
+                    weight = float(jnp.sum(jnp.abs(new_vec) ** 2))
+                    candidates.append(((idx_p << 1) | bit, new_vec, weight))
+            # heapq.nlargest instead of a full sort-then-slice (prog.txt
+            # point 5e): only the top k by weight are ever used below, and
+            # this avoids materializing/sorting the full candidates list
+            # when len(candidates) >> k. Final probability order is
+            # re-derived from scratch at the end of this function anyway
+            # (`order = np.argsort(-probabilities)`), so which of two
+            # equal-weight candidates heapq happens to prefer over sort's
+            # stable order has no effect on the result.
+            paths = [(idx, vec) for idx, vec, _ in heapq.nlargest(k, candidates, key=lambda c: c[2])]
+
+        indices = np.array([p[0] for p in paths])
+        amplitudes = np.array([
+            complex(vec[0]) if len(vec) == 1 else complex(jnp.sum(vec))
+            for _, vec in paths
+        ])
+        probabilities = np.abs(amplitudes) ** 2
+        order = np.argsort(-probabilities)
+        return indices[order], probabilities[order]
+
+    # metrics
+    def max_bond_used(self) -> int:
+        return max(self._bond_history) if self._bond_history else 1
+
+    def total_truncation_error(self) -> float:
+        if not self.truncation_errors:
+            return 0.0
+        return float(np.sqrt(np.sum(np.array(self.truncation_errors) ** 2)))
+
+    def avg_jsd(self) -> float:
+        return float(np.mean(self.jsd_per_bond)) if self.jsd_per_bond else 0.0
+
+    def memory_bytes(self) -> int:
+        bytes_gammas = sum(g.size * g.dtype.itemsize for g in self.gammas)
+        bytes_lambdas = sum(l.size * l.dtype.itemsize for l in self.lambdas)
+        return int(bytes_gammas + bytes_lambdas)
+
+    def memory_mb(self) -> float:
+        return self.memory_bytes() / (1024 * 1024)
+
+    def summary(self) -> str:
+        ee_max = self.entanglement_entropy.max() if len(self.entanglement_entropy) else 0.0
+        return (
+            f"MPSSimulator | n={self.n} | chi_max={self.chi} | "
+            f"chi_used={self.max_bond_used()} | mem={self.memory_mb():.3f}MB | "
+            f"trunc_err={self.total_truncation_error():.2e} | "
+            f"avg_JSD={self.avg_jsd():.4f} | EE_max={ee_max:.3f}b | "
+            f"budget_violations={self.budget_violations}"
+        )
+
+    # ── JIT-fused whole-circuit execution ─────────────────────────────
+    def _record_diag_bookkeeping(self, diag, q1_ids: np.ndarray, is_2q_mask: np.ndarray) -> None:
+        """Shared by both run_circuit_jit paths: jax.lax.scan's stacked
+        per-step diagnostics (diag) replace the eager path's Python
+        list.append()s inside the loop -- same final content, populated
+        differently. Only 2-qubit steps count (is_2q_mask), same as
+        _bond_history/jsd_per_bond only ever growing on 2-qubit gates in
+        the eager path."""
+        chi_history, jsd_history, trunc_err_history, entropy_history = (
+            np.asarray(diag[0]), np.asarray(diag[1]), np.asarray(diag[2]), np.asarray(diag[3]))
+        for i in np.nonzero(is_2q_mask)[0]:
+            chi_new = int(chi_history[i])
+            jsd_val = float(jsd_history[i])
+            self._bond_history.append(chi_new)
+            self.jsd_per_bond.append(jsd_val)
+            self.truncation_errors.append(float(trunc_err_history[i]))
+            q1 = q1_ids[i]
+            if q1 < len(self.entanglement_entropy):
+                self.entanglement_entropy[q1] = float(entropy_history[i])
+            if jsd_val > self.jsd_budget:
+                if self.budget_violations == 0:
+                    warnings.warn(
+                        f"MPSSimulator: bond dimension capped at max_bond={self.chi}, "
+                        f"jsd_budget={self.jsd_budget:.1e} not honored "
+                        f"(jsd={jsd_val:.2e}) -- results may be unreliable, "
+                        f"consider raising max_bond.",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+                self.budget_violations += 1
+
+    def run_circuit_jit(self, ops: List, fuse_gates: bool = False) -> None:
+        """Runs an entire circuit through a single jax.lax.scan-fused,
+        @jax.jit-compiled kernel instead of one eager Python call per gate
+        -- the eager path (apply_gate_1q/apply_gate_2q/_apply_nonlocal_2q,
+        all still available and unchanged) has zero @jax.jit anywhere and
+        pays a host-device sync on every 2-qubit gate's bond-dimension
+        search; measured 88.9s vs Qiskit Aer's 0.64s on a 60-qubit stress
+        circuit -- see README changelog for the real before/after number
+        this method produces on that same circuit.
+
+        Trade-off, explicit and intentional (not hidden): every gamma/
+        lambda is kept at a fixed max_bond-padded size for the rest of
+        this instance's lifetime after this call. Structurally correct
+        either way (zero-padding is mathematically transparent to every
+        other method here -- contract_to_statevector, get_top_k_probable_
+        states, etc. all still work correctly on the padded arrays,
+        verified), just not memory-minimal for genuinely low-entanglement
+        circuits, which is this module's whole point for very large qubit
+        counts. Use the eager methods directly instead when memory, not
+        speed, is the priority -- this is an addition, not a replacement.
+
+        ops: same convention as DenseSVSimulator.run_circuit_jit_beast_mode
+        -- list of (name, *args) tuples/lists. Unlike that method, SWAP is
+        never decomposed into 3xCX (kept as one real gate, see
+        _compile_mps_ops's docstring for why that matters here).
+
+        fuse_gates: opt-in, default False. When True, consecutive gates
+        acting on the same (or a growing) qubit pair are fused into one
+        matrix on the host before compiling (exact -- matrix
+        multiplication, no approximation), cutting the number of scan
+        steps and measurably faster on GPU (~2x on top of the bucketed
+        SVD dispatch alone, see Dense-Evolution-Discovery's
+        mps_gate_blocking_redesign_v2 experiment for the full validation,
+        including verification against non-adjacent-gate and CCX
+        circuits). The trade-off: self._bond_history/jsd_per_bond/
+        truncation_errors/entanglement_entropy get one entry per FUSED
+        step instead of per original gate -- real diagnostics, just
+        coarser-grained. Defaults to False so existing behavior and
+        per-gate bookkeeping granularity are unchanged unless requested.
+        """
+        dtype = self.gammas[0].dtype
+        lambda_dtype = self.lambdas[0].dtype
+
+        if fuse_gates:
+            compiled_rows = _compile_mps_ops(ops, self.n)
+            fused = _fuse_compiled_rows(compiled_rows, dtype) if compiled_rows else []
+
+            if self._fused_mps_runner is None:
+                self._fused_mps_runner = _build_fused_mps_runner(self.n, self.chi, self.eps, self.jsd_budget)
+
+            gammas_padded = _pad_all_gammas(tuple(self.gammas), self.chi, dtype)
+            lambdas_padded = _pad_all_lambdas(tuple(self.lambdas), self.chi, lambda_dtype)
+            real_chi_initial = jnp.asarray(self._real_chi, dtype=jnp.int32)
+
+            if fused:
+                xs = _fused_entries_to_arrays(fused, dtype)
+                final_gammas, final_lambdas, final_real_chi, diag = self._fused_mps_runner(
+                    gammas_padded, lambdas_padded, real_chi_initial, xs)
+
+                self.gammas = [final_gammas[i] for i in range(self.n)]
+                self.lambdas = [final_lambdas[i] for i in range(self.n + 1)]
+                self._real_chi = np.asarray(final_real_chi)
+
+                q1_ids = np.asarray([entry[1] for entry in fused])
+                is_2q_mask = np.asarray([entry[0] == '2q' for entry in fused])
+                self._record_diag_bookkeeping(diag, q1_ids, is_2q_mask)
+            return
+
+        compiled_rows = _compile_mps_ops(ops, self.n)
+        ops_dtype = _real_dtype_for(dtype)
+
+        if compiled_rows:
+            ops_array = jnp.array(compiled_rows, dtype=ops_dtype)
+        else:
+            ops_array = jnp.zeros((0, 5), dtype=ops_dtype)
+
+        if self._mps_runner is None:
+            self._mps_runner = _build_mps_runner(self.n, self.chi, self.eps, self.jsd_budget)
+
+        gammas_padded = jnp.stack([_pad_gamma(g, self.chi).astype(dtype) for g in self.gammas])
+        lambdas_padded = jnp.stack([_pad_lambda(l, self.chi).astype(lambda_dtype) for l in self.lambdas])
+        real_chi_initial = jnp.asarray(self._real_chi, dtype=jnp.int32)
+
+        final_gammas, final_lambdas, final_real_chi, diag = self._mps_runner(
+            gammas_padded, lambdas_padded, real_chi_initial, ops_array)
+
+        self.gammas = [final_gammas[i] for i in range(self.n)]
+        self.lambdas = [final_lambdas[i] for i in range(self.n + 1)]
+        self._real_chi = np.asarray(final_real_chi)
+
+        if compiled_rows:
+            g_ids = np.asarray([row[0] for row in compiled_rows])
+            q1_ids = np.asarray([int(row[1]) for row in compiled_rows])
+            is_2q_mask = g_ids >= 20
+            self._record_diag_bookkeeping(diag, q1_ids, is_2q_mask)
+
+
+_PAULI_MATRICES = {
+    'I': np.eye(2, dtype=np.complex128),
+    'X': np.array([[0, 1], [1, 0]], dtype=np.complex128),
+    'Y': np.array([[0, -1j], [1j, 0]], dtype=np.complex128),
+    'Z': np.array([[1, 0], [0, -1]], dtype=np.complex128),
+}
+
+
+def _mps_transfer_sweep(
+    mps: "MPSSimulator", assignment: dict, need_norm: bool = True
+) -> Tuple[complex, Optional[float]]:
+    """Left-to-right transfer-matrix sweep over the Gamma/Lambda tensors
+    computing <psi|P|psi> (unnormalized) for the per-site operator
+    assignment (site -> 'I'|'X'|'Y'|'Z', missing sites default to 'I') --
+    never materializes a (2**n,) statevector, cost O(n * chi^3). When
+    need_norm is True, <psi|psi> is accumulated in the same Python loop
+    (reusing the per-site A = Lambda*Gamma tensor and its conjugate) and
+    returned alongside the raw value, instead of a second, separate sweep
+    -- when False (used by mps_pauli_sum_expectation, which only needs
+    the norm once for the whole sum, not once per term) that accumulation
+    is skipped entirely.
+
+    Every site (not only ones in the assignment) goes through the
+    identical per-site contraction, using the identity matrix wherever the
+    assignment doesn't specify a Pauli -- this is already O(chi^3) per
+    site regardless of whether the local operator is I or a real Pauli, so
+    there is no separate "skip identity sites" fast path to get right or
+    wrong: the same sweep is correct for any placement of the operator's
+    support, including support that spans much of the chain, without
+    relying on which side of any notional orthogonality center a given
+    site falls on.
+
+    The norm is needed at all because SVD truncation (apply_gate_2q) can
+    leave the Gamma/Lambda tensors at less than exact unit norm --
+    contract_to_statevector renormalizes explicitly for the same reason.
+    If the MPS ever carried exact unit norm this division is a no-op, but
+    every reader of the state normalizes explicitly rather than assuming
+    that invariant holds.
+    """
+    dtype = mps.gammas[0].dtype
+    env_op = jnp.ones((1, 1), dtype=dtype)
+    env_norm = jnp.ones((1, 1), dtype=dtype) if need_norm else None
+    for i in range(mps.n):
+        op = _PAULI_MATRICES[assignment.get(i, 'I')].astype(dtype)
+        a = jnp.einsum('l,lpr->lpr', mps.lambdas[i].astype(dtype), mps.gammas[i])
+        a_conj = jnp.conj(a)
+        a_op = jnp.einsum('pq,lqr->lpr', op, a)
+        env_op = jnp.einsum('lk,lpr->kpr', env_op, a_conj)
+        env_op = jnp.einsum('kpr,kps->rs', env_op, a_op)
+        if need_norm:
+            env_norm = jnp.einsum('lk,lpr->kpr', env_norm, a_conj)
+            env_norm = jnp.einsum('kpr,kps->rs', env_norm, a)
+    raw = complex(env_op[0, 0])
+    norm_sq = complex(env_norm[0, 0]).real if need_norm else None
+    return raw, norm_sq
+
+
+def mps_pauli_expectation(mps: "MPSSimulator", pauli_terms) -> complex:
+    """<psi|P|psi> / <psi|psi> for a single Pauli string P, contracted
+    directly against the MPS (Gamma/Lambda tensors).
+
+    pauli_terms accepts the same three forms as
+    physics.observables.pauli_expectation (a string, e.g. 'XIZ'; a dict
+    {qubit: 'X'|'Y'|'Z'}; or an iterable of (qubit, pauli) pairs) -- reuses
+    that module's own `_normalize_terms` so both functions agree on
+    parsing by construction, not by parallel reimplementation. See
+    _mps_transfer_sweep for why the division is needed.
+    """
+    assignment = _normalize_terms(pauli_terms, mps.n)
+    raw, norm_sq = _mps_transfer_sweep(mps, assignment, need_norm=True)
+    return raw / norm_sq
+
+
+def mps_pauli_sum_expectation(mps: "MPSSimulator", terms) -> complex:
+    """sum_i coeff_i * <psi|P_i|psi> / <psi|psi> -- same terms format as
+    physics.observables.pauli_sum_expectation: an iterable of
+    (coeff, pauli_terms) pairs. <psi|psi> does not depend on which Pauli
+    string is being measured, so it is computed once via its own sweep
+    and applied to the whole sum, instead of once per term."""
+    terms = list(terms)
+    if not terms:
+        return 0j
+    _, norm_sq = _mps_transfer_sweep(mps, {}, need_norm=True)
+    raw_sum = sum(
+        coeff * _mps_transfer_sweep(mps, _normalize_terms(pauli_terms, mps.n), need_norm=False)[0]
+        for coeff, pauli_terms in terms
+    )
+    return raw_sum / norm_sq
+
+
+@dataclasses.dataclass
+class BondConvergenceResult:
+    bonds: List[int]
+    chi_used: List[int]
+    avg_jsd: List[float]
+    budget_violations: List[int]
+    values: List[List[complex]]
+    diffs: List[List[float]]
+    verdicts: List[str]
+
+
+def bond_convergence(
+    ops: List, n_qubits: int, observables: list, bonds: List[int],
+    tol: float = 1e-3, **mps_kwargs,
+) -> BondConvergenceResult:
+    """Runs the same circuit at every value in `bonds` (increasing) and
+    checks whether the reported observables have actually converged with
+    respect to bond dimension, instead of trusting a single run's own
+    internal diagnostics.
+
+    Requires len(bonds) >= 3. Two bonds give exactly one discrepancy,
+    which is a single number with no way to tell whether it is still
+    shrinking toward `tol` or has already stalled -- measured on a
+    40-qubit, 4-layer brickwall circuit, chi=4->8->32 gave |<Z0>|
+    discrepancies of ~4.7e-2 then ~1.2e-2 (chi_used never hit its own
+    cap, so this is a real not_converged, not an artifact of running out
+    of bond dimension): a two-bond check (chi=4 vs 8) would see only the
+    first number and have no basis to call it anything, while three
+    bonds show a trend that is decreasing but still two orders of
+    magnitude above any reasonable `tol`.
+
+    A verdict of "converged" additionally requires the successive
+    discrepancies to be monotonically non-increasing, not just that the
+    last one is below `tol` -- a single small discrepancy proves nothing
+    about the trend on its own, which is the same failure mode as the
+    two-bond case above, one level up. (Ties count as non-increasing: an
+    exactly-converged observable, e.g. a GHZ chain whose bond dimension
+    never needs to grow, produces identical values -- and therefore
+    zero discrepancies -- at every bond, which must count as converged.)
+
+    avg_jsd and budget_violations (from the underlying MPSSimulator runs)
+    are reported per bond for context only, never used to decide the
+    verdict -- a low average JSD is computed per truncation step and says
+    nothing about whether the specific observable being tracked has
+    settled down as `max_bond` grows.
+
+    If max_bond_used() at the highest bond still equals that bond's cap,
+    the truncation never had headroom below max_bond at any cut, so no
+    tolerance can be certified from this data: every observable's verdict
+    becomes "undecidable" regardless of its own discrepancies.
+    """
+    if len(bonds) < 3:
+        raise ValueError(f"bond_convergence needs at least 3 bonds to detect a trend, got {len(bonds)}")
+
+    chi_used, avg_jsd, budget_violations = [], [], []
+    values = [[] for _ in observables]
+    for bond in bonds:
+        mps = MPSSimulator(n_qubits=n_qubits, max_bond=bond, **mps_kwargs)
+        mps.run_circuit_jit(ops)
+        chi_used.append(mps.max_bond_used())
+        avg_jsd.append(mps.avg_jsd())
+        budget_violations.append(mps.budget_violations)
+        for obs_idx, obs in enumerate(observables):
+            values[obs_idx].append(mps_pauli_expectation(mps, obs))
+
+    diffs = [
+        [abs(vals[i + 1] - vals[i]) for i in range(len(vals) - 1)]
+        for vals in values
+    ]
+
+    undecidable = chi_used[-1] >= bonds[-1]
+    verdicts = []
+    for d in diffs:
+        if undecidable:
+            verdicts.append("undecidable")
+        elif all(d[i + 1] <= d[i] for i in range(len(d) - 1)) and d[-1] < tol:
+            verdicts.append("converged")
+        else:
+            verdicts.append("not_converged")
+
+    return BondConvergenceResult(
+        bonds=list(bonds), chi_used=chi_used, avg_jsd=avg_jsd,
+        budget_violations=budget_violations, values=values, diffs=diffs,
+        verdicts=verdicts,
+    )