Measure MPS norms with tn_norm

Author: rezaquant

history/2026-06-24-dimension-aware-swap-routing.md

@@ -1,7 +1,7 @@
 # 2026-06-24 — Dimension-aware SWAP routing
 
 - Milestone: B0 — gate-routing audit / quimb integration prerequisite
-- Branch / commit: `main` working tree, uncommitted
+- Branch / commit: `main` / `2caae41`
 
 ## What changed
 - Added dimension-aware routed SWAP construction in `src/pepsy/operators/gates.py`.

history/2026-06-24-mps-tn-norm-normalization.md

@@ -0,0 +1,39 @@
+# 2026-06-24 — MPS normalization via `tn_norm`
+
+- Milestone: M1 — non-unitary MPS diagnostics and normalization hardening
+- Branch / commit: `main` working tree, pending commit
+
+## What changed
+- Switched `MpsOptimizer._canonical_span_norm(...)` from contracting a selected
+  canonical span into a dense tensor to measuring the raw working-data norm with
+  `pepsy.tensors.core.tn_norm(..., strip_exponent=True)`.
+- Temporarily clears and restores `p.exponent` while measuring the raw tensor
+  data, so automatic normalization and `track_norm_infidelity=True` do not
+  accidentally include the represented global scale.
+- Expanded non-unitary norm-infidelity smoke coverage across `dmrg`, `mpo`,
+  `swap`, and `svd` modes, with an explicit skip for older quimb versions that
+  do not expose `gate_with_auto_swap_`.
+
+## Why
+- The old span contraction was correct only when the canonical span stayed
+  small. For long-range gates or swap/sweep paths, the span can retain many
+  physical legs; contracting it as one dense block can explode even though the
+  norm itself is just a scalar double-layer contraction.
+- `tn_norm(..., strip_exponent=True)` keeps the measurement consistent with the
+  Pepsy normalization model: normalize raw data locally, store the removed scale
+  in `p.exponent`, and leave `p.norm()` to report the represented state norm.
+
+## How it was validated
+- `python -m pytest -q tests/test_optimize_mps.py` -> `48 passed`.
+
+## Decisions / findings
+- The `fallback` argument remains accepted for compatibility, but the measured
+  path now uses `tn_norm` directly rather than the old dense-span fallback.
+- True overlap fidelity still goes through `tn_fidelity`; this change only
+  affects raw norm measurements used by normalization and norm-infidelity
+  diagnostics.
+
+## Next step
+- Keep the Tensy `to_mps(..., contraction_opt=optimizer_cotengra)` path wired
+  through this same contraction optimizer so large DEM/PF non-unitary streams
+  use the intended double-layer contraction route for norm diagnostics.

src/pepsy/optimizers/mps/optimizer.py

@@ -25,7 +25,7 @@
 
 from ...fitting.local import FIT
 from ...operators.gates import _normalize_gate_entries, gate as apply_gate
-from ...tensors.core import tn_fidelity
+from ...tensors.core import tn_fidelity, tn_norm
 
 __all__ = ["MpsOptimizer"]
 
@@ -69,7 +69,7 @@ class MpsOptimizer:  # pylint: disable=too-many-instance-attributes
     ----------
     normalizations : list[dict]
         Automatic normalization events recorded during :meth:`run`. Each entry
-        stores the 1-based gate step, previous local squared norm,
+        stores the 1-based gate step, previous raw squared norm,
         canonicalization span, tensor site where the normalization factor was
         inserted, and resulting base-10 ``p.exponent``.
         The raw tensor data are rescaled; the represented norm remains
@@ -637,28 +637,29 @@ def _normalize_span(where):
         raise ValueError("where must be an int, (int,), or (int, int).")
 
     def _canonical_span_norm(self, p, where, *, fallback=True):
-        """Return ``||p||`` from a canonical center/span block.
+        """Return the raw working-data norm without densifying wide spans.
 
-        This assumes tensors outside ``where`` are already isometric. The span
-        itself can contain multiple tensors per site, e.g. after ``split-gate``.
+        Normalization needs the norm of the current tensor data, excluding the
+        accumulated ``p.exponent`` scale. Contracting a canonical span into a
+        dense block can explode for long-range gates because the span retains
+        all physical legs, so use ``tn_norm``'s double-layer contraction instead.
         """
-        xmin, xmax = self._normalize_span(where)
+        _ = where, fallback
+        exponent = getattr(p, "exponent", None)
         try:
-            tags = [p.site_tag(i) for i in range(xmin, xmax + 1)]
-            block = p.select(tags, which="any")
-            if isinstance(block, qtn.TensorNetwork):
-                if block.num_tensors == 0:
-                    raise ValueError("canonical span selected no tensors.")
-                block = block.contract(
-                    all,
-                    output_inds=block.outer_inds(),
-                    optimize=self.contraction_opt,
-                )
-            return ar.do("linalg.norm", block.data)
-        except Exception:
-            if not fallback:
-                raise
-            return p.norm(optimize=self.contraction_opt)
+            if exponent is not None:
+                p.exponent = 0.0
+            mantissa, exponent_sq = tn_norm(
+                p,
+                contraction_opt=self.contraction_opt,
+                strip_exponent=True,
+            )
+            return ar.do("sqrt", ar.do("abs", mantissa)) * 10 ** (
+                float(exponent_sq) / 2.0
+            )
+        finally:
+            if exponent is not None:
+                p.exponent = exponent
 
     @staticmethod
     def _norm_ratio_fidelity(approx_norm, target_norm):

tests/test_optimize_mps.py

@@ -10,6 +10,7 @@
 import quimb.tensor as qtn
 
 import pepsy as py
+import pepsy.optimizers.mps.optimizer as mps_optimizer_module
 
 
 def _non_unitary_entangling_gate():
@@ -386,6 +387,54 @@ def test_mps_optimizer_canonical_span_norm_matches_full_target_norm(where):
     assert local_norm == pytest.approx(target.norm())
 
 
+def test_mps_optimizer_canonical_span_norm_ignores_stored_exponent():
+    """Internal normalization should measure raw data, not represented scale."""
+    p0 = qtn.MPS_rand_state(4, bond_dim=2, phys_dim=2, dtype="complex128")
+    opt = py.MpsOptimizer(p0.copy(), gates=[], chi=8, mode="svd")
+    opt.p.exponent = 3.0
+
+    raw = opt.p.copy()
+    raw.exponent = 0.0
+    measured = opt._canonical_span_norm(opt.p, (0, 3))  # pylint: disable=protected-access
+
+    assert measured == pytest.approx(raw.norm())
+    assert opt.p.exponent == pytest.approx(3.0)
+
+
+def test_mps_optimizer_norm_infidelity_uses_tn_norm_strip_exponent(monkeypatch):
+    """Norm-infidelity diagnostics should measure raw norms through ``tn_norm``."""
+    calls = []
+    original_tn_norm = mps_optimizer_module.tn_norm
+
+    def _spy_tn_norm(*args, **kwargs):
+        calls.append(kwargs.copy())
+        return original_tn_norm(*args, **kwargs)
+
+    monkeypatch.setattr(mps_optimizer_module, "tn_norm", _spy_tn_norm)
+
+    p0 = qtn.MPS_computational_state("0000", dtype="complex128")
+    gates = [
+        (qu.hadamard(), (0,)),
+        (qu.hadamard(), (1,)),
+        (_non_unitary_entangling_gate(), (0, 1)),
+    ]
+
+    opt = py.MpsOptimizer(p0.copy(), gates=gates, chi=1, mode="mpo")
+    opt.run(
+        progbar=False,
+        cutoff=1e-12,
+        non_unitary=True,
+        normalize_final=True,
+        track_norm_infidelity=True,
+    )
+
+    samples = opt.get_norm_infidelity_samples()
+    assert len(samples) == 1
+    assert len(calls) >= 2
+    assert all(call["strip_exponent"] is True for call in calls)
+    assert all(call["contraction_opt"] == opt.contraction_opt for call in calls)
+
+
 def test_mps_optimizer_non_unitary_norm_infidelity_matches_svd_target():
     """SVD non-unitary proxy should match quimb's target infidelity."""
     p0 = qtn.MPS_computational_state("0000", dtype="complex128")
@@ -664,9 +713,9 @@ def test_mps_optimizer_dmrg_non_unitary_matches_mpo_accuracy():
     )
 
 
-@pytest.mark.parametrize("mode", ["dmrg", "mpo"])
+@pytest.mark.parametrize("mode", ["dmrg", "mpo", "swap", "svd"])
 def test_mps_optimizer_non_unitary_norm_infidelity_smoke_other_modes(mode):
-    """Other compressed modes should expose a bounded non-unitary proxy."""
+    """All compressed modes should expose a bounded non-unitary proxy."""
     p0 = qtn.MPS_computational_state("0000", dtype="complex128")
     gates = [
         (qu.hadamard(), (0,)),
@@ -675,6 +724,8 @@ def test_mps_optimizer_non_unitary_norm_infidelity_smoke_other_modes(mode):
     ]
 
     opt = py.MpsOptimizer(p0.copy(), gates=gates, chi=1, mode=mode)
+    if mode == "swap" and not hasattr(opt.p, "gate_with_auto_swap_"):
+        pytest.skip("swap mode requires gate_with_auto_swap_ in this quimb version.")
     opt.run(
         progbar=False,
         cutoff=1e-12,