pdft 0.2.1__tar.gz → 0.2.2__tar.gz

{pdft-0.2.1 → pdft-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pdft
-Version: 0.2.1
+Version: 0.2.2
 Summary: Python port of ParametricDFT.jl: learning parametric quantum Fourier transforms via manifold optimization.
 Project-URL: Homepage, https://github.com/zazabap/pdft
 Project-URL: Issues, https://github.com/zazabap/pdft/issues

{pdft-0.2.1 → pdft-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pdft"
-version = "0.2.1"
+version = "0.2.2"
 description = "Python port of ParametricDFT.jl: learning parametric quantum Fourier transforms via manifold optimization."
 readme = "README.md"
 requires-python = ">=3.11"

{pdft-0.2.1 → pdft-0.2.2}/src/pdft/__init__.py

@@ -41,7 +41,7 @@ if not _os.environ.get("PDFT_DISABLE_COMPILE_CACHE"):
     except OSError:
         pass
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"
 __upstream_ref__ = "nzy1997/ParametricDFT.jl@a201a27e47df2f0f3ab460f83d49b6e5f5d1e9ef"
 
 # Slim public re-export hub: most-used names only. Anything else: import

{pdft-0.2.1 → pdft-0.2.2}/src/pdft/circuit/builder.py

@@ -167,6 +167,92 @@ def build_circuit_einsum(
     return subscripts, tensor_list, tensor_shapes
 
 
+def _axis_of_qubit(q: int, m: int, n: int) -> int:
+    """Map a 1-indexed qubit number to its position in the (2,)*N reshape of pic.
+
+    Yao's little-endian convention (matching `build_circuit_einsum`'s
+    `pic_labels` ordering): qubit `m` maps to axis 0, qubit `m-1` to axis 1,
+    …, qubit 1 to axis m-1, qubit `m+n` to axis m, …, qubit `m+1` to axis
+    m+n-1.
+    """
+    if 1 <= q <= m:
+        return m - q
+    if m + 1 <= q <= m + n:
+        return m + (m + n - q)
+    raise ValueError(f"qubit index {q} out of range (1..{m + n})")
+
+
+def _stepped_apply(
+    program: list[tuple[str, tuple[int, ...]]],
+    tensor_index: list[int],
+    operands: tuple,
+    *,
+    m: int,
+    n: int,
+    inverse: bool,
+) -> Array:
+    """Apply gate program to pic via per-gate `jnp.tensordot` calls.
+
+    operands = (*tensors, pic) where pic has shape (2,) * (m+n) and tensors
+    is the operand list passed to the closure (sorted Hadamard-first by
+    `compile_circuit`, matching the convention of the legacy big-einsum
+    path so saved checkpoints load with the same axis order).
+
+    `tensor_index[i]` maps program step i (in temporal order) to its index
+    in the operand `tensors` tuple.
+    """
+    *tensors, pic = operands
+    for i, (kind, qubits) in enumerate(program):
+        T = tensors[tensor_index[i]]
+        if kind == "H":
+            (q,) = qubits
+            ax = _axis_of_qubit(q, m, n)
+            # Forward: H[out, in] · pic[..., ax=in, ...] -> pic[..., ax=out, ...]
+            #   contract on T's axis 1 (the "in" leg).
+            # Inverse: apply the transpose of T. Combined with the caller's
+            #   conj(T) (loss._scalar_loss conjugates tensors before passing
+            #   them to the inverse closure), this realises the true adjoint
+            #   T† = conj(T.T). For symmetric H this collapses to conj(T),
+            #   but for trained / non-symmetric H tensors the leg-swap is
+            #   essential — without it the round-trip T† T x ≠ x. The U4
+            #   handler below uses the same forward/inverse axis convention.
+            contract_in = 0 if inverse else 1
+            pic = jnp.tensordot(T, pic, axes=[[contract_in], [ax]])
+            pic = jnp.moveaxis(pic, 0, ax)
+        elif kind == "CP":
+            q_ctrl, q_tgt = qubits
+            ax_c = _axis_of_qubit(q_ctrl, m, n)
+            ax_t = _axis_of_qubit(q_tgt, m, n)
+            # T[c, t] is the diagonal of CP. pic *= T broadcast onto (ax_c, ax_t).
+            broadcast_shape = [1] * pic.ndim
+            broadcast_shape[ax_c] = 2
+            broadcast_shape[ax_t] = 2
+            # T's axis 0 should land at ax_c, axis 1 at ax_t. If ax_c < ax_t,
+            # the natural reshape works; if ax_c > ax_t, we transpose first
+            # so the smaller dimension index sees T's first axis.
+            T_oriented = T if ax_c < ax_t else T.T
+            pic = pic * T_oriented.reshape(broadcast_shape)
+        elif kind == "U4":
+            q_ctrl, q_tgt = qubits
+            ax_c = _axis_of_qubit(q_ctrl, m, n)
+            ax_t = _axis_of_qubit(q_tgt, m, n)
+            # Forward: T[oc, ot, ic, it] · pic[..., ax_c=ic, ax_t=it, ...]
+            #          -> pic[..., ax_c=oc, ax_t=ot, ...]
+            # Inverse (transpose, not adjoint): T[oc, ot, ic, it]
+            #          · pic[..., ax_c=oc, ax_t=ot, ...]
+            #          -> pic[..., ax_c=ic, ax_t=it, ...]
+            # The conjugate flip for the true unitary inverse is applied by
+            # the caller (loss._scalar_loss conjugates tensors before
+            # calling the inverse closure), matching the legacy
+            # build_circuit_einsum behaviour.
+            contract_axes = [[0, 1], [ax_c, ax_t]] if inverse else [[2, 3], [ax_c, ax_t]]
+            pic = jnp.tensordot(T, pic, axes=contract_axes)
+            pic = jnp.moveaxis(pic, [0, 1], [ax_c, ax_t])
+        else:
+            raise AssertionError(f"unknown gate kind: {kind}")
+    return pic
+
+
 def compile_circuit(
     gates: list[Gate],
     m: int,
@@ -174,12 +260,61 @@ def compile_circuit(
     *,
     inverse: bool,
 ) -> tuple[Callable[..., Array], list[Array]]:
-    """High-level entry: build subscripts and jit-compile the einsum."""
-    from .cache import optimize_code_cached
+    """Build a JIT-compiled stepped circuit closure.
+
+    Returns ``(code, tensors)`` where ``code(*tensors, pic_reshaped)``
+    applies the gates one at a time via :func:`jnp.tensordot`, recycling
+    contraction labels across steps so the 52-character a-zA-Z label pool
+    of the legacy single-einsum path is never reached. Each step consumes
+    at most ~22 wire labels regardless of circuit size.
+
+    The returned ``tensors`` list preserves the Hadamard-first sort applied
+    by the legacy builder, so saved trained checkpoints (which serialize
+    tensors in this order) continue to load consistently.
+    """
+    n_gates = len(gates)
+    program: list[tuple[str, tuple[int, ...]]] = [
+        (g["kind"], g["qubits"]) for g in gates
+    ]
+    tensor_list: list[Array] = [g["tensor"] for g in gates]
+
+    # Hadamard-first sort — preserves the legacy convention for the
+    # returned tensor list. `perm[i]` is the original (temporal) gate index
+    # of the i-th sorted slot, so `sorted_tensors[i] = original[perm[i]]`.
+    import numpy as _np
+
+    H_np = _np.asarray(HADAMARD)
+
+    def _is_hadamard(t):
+        a = _np.asarray(t)
+        return a.shape == (2, 2) and _np.allclose(a, H_np, atol=1e-12)
+
+    is_not_hadamard = [not _is_hadamard(t) for t in tensor_list]
+    perm = sorted(range(n_gates), key=lambda i: is_not_hadamard[i])
+    sorted_tensors = [tensor_list[i] for i in perm]
+
+    # Map temporal gate index → position in the sorted operand tuple.
+    # `temporal_to_sorted[j] = i` such that perm[i] = j.
+    temporal_to_sorted = [0] * n_gates
+    for sorted_pos, original_idx in enumerate(perm):
+        temporal_to_sorted[original_idx] = sorted_pos
+
+    if inverse:
+        # Apply gates in reverse temporal order. tensor_index aligns with
+        # the reversed program: program_step[i] needs the tensor for the
+        # ORIGINAL gate at index (n_gates-1-i).
+        program = list(reversed(program))
+        tensor_index = [temporal_to_sorted[n_gates - 1 - i] for i in range(n_gates)]
+    else:
+        tensor_index = [temporal_to_sorted[i] for i in range(n_gates)]
+
+    @jax.jit
+    def code(*operands):
+        return _stepped_apply(
+            program, tensor_index, operands, m=m, n=n, inverse=inverse,
+        )
 
-    subscripts, tensors, shapes = build_circuit_einsum(gates, m, n, inverse=inverse)
-    code = optimize_code_cached(subscripts, *shapes)
-    return code, tensors
+    return code, sorted_tensors
 
 
 def apply_circuit(

{pdft-0.2.1 → pdft-0.2.2}/src/pdft/loss.py

@@ -120,7 +120,13 @@ def _scalar_loss(
         truncated = topk_truncate(pred, loss.k)
         conj_tensors = [jnp.conj(t) for t in tensors]  # type: ignore[arg-type]
         reconstructed = _apply_circuit(conj_tensors, inverse_code, m, n, truncated)  # type: ignore[arg-type]
-        return jnp.sum(jnp.abs(target - reconstructed) ** 2)
+        base = jnp.sum(jnp.abs(target - reconstructed) ** 2)
+        # Optional hook for MSELoss subclasses that add regularizers or
+        # auxiliary scalar terms while preserving vanilla MSELoss behavior.
+        extra_fn = getattr(loss, "_extra_loss", None)
+        if extra_fn is not None:
+            return base + extra_fn(tensors)
+        return base
     raise TypeError(f"unsupported loss type: {type(loss).__name__}")
 
 

{pdft-0.2.1 → pdft-0.2.2}/src/pdft/optimizers/core.py

@@ -48,13 +48,26 @@ def _common_setup(tensors: list[Array]) -> _OptimizationState:
     )
 
 
-def _batched_project(state: _OptimizationState, euclid_grads: list[Array]):
+def _batched_project(
+    state: _OptimizationState,
+    euclid_grads: list[Array],
+    frozen_indices: frozenset[int] | None = None,
+):
     """Mirror of upstream src/optimizers.jl:129-149."""
     rg_batches: dict[AbstractRiemannianManifold, Array] = {}
     grad_norm_sq = 0.0
     for manifold, indices in state.manifold_groups.items():
         pb = state.point_batches[manifold]
-        gb = stack_tensors(euclid_grads, indices)
+        if frozen_indices:
+            gb = jnp.stack(
+                [
+                    jnp.zeros_like(euclid_grads[i]) if i in frozen_indices else euclid_grads[i]
+                    for i in indices
+                ],
+                axis=-1,
+            )
+        else:
+            gb = stack_tensors(euclid_grads, indices)
         rg = manifold.project(pb, gb)
         rg_batches[manifold] = rg
         grad_norm_sq = grad_norm_sq + float(jnp.real(jnp.sum(jnp.conj(rg) * rg)))

{pdft-0.2.1 → pdft-0.2.2}/src/pdft/optimizers/loop.py

@@ -25,6 +25,7 @@ def optimize(
     max_iter: int = 100,
     tol: float = 1e-6,
     record_loss: bool = False,
+    frozen_indices: frozenset[int] | None = None,
 ) -> tuple[list[Array], list[float]]:
     """Mirror of upstream src/optimizers.jl:335-412.
 
@@ -61,7 +62,7 @@ def optimize(
                 warnings.warn("Non-finite gradient — optimizer stopping.", stacklevel=2)
                 return state.current_tensors, trace
 
-        rg_batches, grad_norm = _batched_project(state, raw_grads)
+        rg_batches, grad_norm = _batched_project(state, raw_grads, frozen_indices=frozen_indices)
         grad_norm_sq = float(grad_norm) ** 2
 
         if opt.max_grad_norm is not None and float(grad_norm) > opt.max_grad_norm:

{pdft-0.2.1 → pdft-0.2.2}/src/pdft/training/adam_step.py

@@ -25,6 +25,7 @@ def _build_jit_adam_step(
     beta2: float,
     eps: float,
     max_grad_norm: float | None,
+    frozen_set: frozenset[int] | None = None,
 ):
     """Build a single JIT'd Adam step for `train_basis_batched`'s fast path.
 
@@ -71,6 +72,17 @@ def _build_jit_adam_step(
         else:
             ibs.append(None)
 
+    # Pre-compute per-manifold-group, per-slot frozen masks (static booleans,
+    # closed over in the JIT'd step_fn — no recompiles because they're Python).
+    # slot_frozen_masks[k] is a tuple of bools, one per slot in manifold group k.
+    if frozen_set:
+        slot_frozen_masks = [
+            tuple(idx in frozen_set for idx in idxs)
+            for idxs in indices_list
+        ]
+    else:
+        slot_frozen_masks = [None] * len(indices_list)
+
     @jax.jit
     def step_fn(tensors_list, m_list, v_list, batch, lr, iter_1based):
         # Forward + backward; loss comes "for free" alongside grads.
@@ -82,9 +94,19 @@ def _build_jit_adam_step(
         # Per-manifold project (fused into the JIT'd graph).
         pb_list = []
         rg_list = []
-        for manifold, idxs in zip(manifold_list, indices_list):
+        for k, (manifold, idxs) in enumerate(zip(manifold_list, indices_list)):
             pb = jnp.stack([tensors_list[i] for i in idxs], axis=-1)
-            gb = jnp.stack([grads[i] for i in idxs], axis=-1)
+            slot_mask = slot_frozen_masks[k]
+            if slot_mask is not None and any(slot_mask):
+                # Zero gradient for frozen slots before projection so Adam's
+                # moment buffers stay zero and the tensors don't move.
+                gb = jnp.stack(
+                    [jnp.zeros_like(grads[i]) if slot_mask[k2] else grads[i]
+                     for k2, i in enumerate(idxs)],
+                    axis=-1,
+                )
+            else:
+                gb = jnp.stack([grads[i] for i in idxs], axis=-1)
             rg = manifold.project(pb, gb)
             pb_list.append(pb)
             rg_list.append(rg)
@@ -120,11 +142,16 @@ def _build_jit_adam_step(
 
             new_m_list.append(new_m)
             new_v_list.append(new_v)
+            slot_mask = slot_frozen_masks[k]
             for k2, idx in enumerate(idxs):
                 # Use ellipsis so this works for any tensor storage shape:
                 # (2, 2, n) for 2x2 unitaries, (2, 2, 2, 2, n) for 2-qubit
                 # (Unitary2qManifold) gates, etc.
-                new_tensors[idx] = new_pb[..., k2]
+                if slot_mask is not None and slot_mask[k2]:
+                    # Frozen: restore the original tensor, bit-exactly.
+                    new_tensors[idx] = tensors_list[idx]
+                else:
+                    new_tensors[idx] = new_pb[..., k2]
 
         return new_tensors, new_m_list, new_v_list, loss_val
 

{pdft-0.2.1 → pdft-0.2.2}/src/pdft/training/batched.py

@@ -13,6 +13,7 @@ The eval+early-stopping bookkeeping is shared via training.eval_loop.
 from __future__ import annotations
 
 import math
+import operator
 import time
 from collections.abc import Sequence
 
@@ -91,6 +92,47 @@ def _validate_batched_args(
         raise ValueError(f"warmup_frac must be in [0, 1), got {warmup_frac}")
 
 
+def _validate_frozen_indices(frozen_indices: "list[int] | None", n_tensors: int) -> frozenset:
+    """Validate and normalise ``frozen_indices``.
+
+    Returns a ``frozenset[int]`` of validated frozen indices (empty set means
+    no freezing).  Raises ``ValueError`` on any violation.
+    """
+    if frozen_indices is None or len(frozen_indices) == 0:
+        return frozenset()
+    seen: set[int] = set()
+    for raw_i in frozen_indices:
+        if isinstance(raw_i, bool):
+            raise ValueError(
+                f"frozen_indices contains non-integer index {raw_i!r}; "
+                "all indices must be integers."
+            )
+        try:
+            i = operator.index(raw_i)
+        except TypeError as exc:
+            raise ValueError(
+                f"frozen_indices contains non-integer index {raw_i!r}; "
+                "all indices must be integers."
+            ) from exc
+        if i < 0:
+            raise ValueError(
+                f"frozen_indices contains negative index {i}; "
+                f"all indices must be in [0, {n_tensors - 1}]."
+            )
+        if i >= n_tensors:
+            raise ValueError(
+                f"frozen_indices contains out-of-range index {i}; "
+                f"basis has {n_tensors} tensors (valid range 0..{n_tensors - 1})."
+            )
+        if i in seen:
+            raise ValueError(
+                f"frozen_indices contains duplicate index {i}; "
+                "each index must appear at most once."
+            )
+        seen.add(i)
+    return frozenset(seen)
+
+
 def train_basis_batched(
     basis,
     *,
@@ -108,12 +150,29 @@ def train_basis_batched(
     shuffle: bool = True,
     seed: int = 0,
     val_every_k_epochs: int = 1,
+    frozen_indices: "list[int] | None" = None,
 ) -> TrainingResult:
     """Multi-image, multi-epoch trainer with cosine LR schedule.
 
     Mirror of `ParametricDFT.jl/src/training.jl::_train_basis_core` (main).
     See parameter docs in the original training.py docstring.
+
+    Parameters
+    ----------
+    frozen_indices : list[int] | None, optional
+        List of integer indices into ``basis.tensors``.  Tensors at these
+        indices are NOT updated during training — they stay at their initial
+        values throughout.  Each step computes gradients on all tensors
+        normally; the update is then suppressed for frozen indices BEFORE any
+        optimizer state is mutated (so Adam's moment buffers for frozen indices
+        remain zero).  Useful for experiments that train only a subset of a
+        circuit's gates.
+
+        Validation: all indices must satisfy ``0 <= i < len(basis.tensors)``,
+        no duplicates are allowed, and an empty list is treated as ``None``
+        (no-op).  ``ValueError`` is raised on mis-specification.
     """
+    frozen_set = _validate_frozen_indices(frozen_indices, len(list(basis.tensors)))
     _validate_batched_args(
         dataset, epochs, batch_size, validation_split, early_stopping_patience, warmup_frac
     )
@@ -188,6 +247,7 @@ def train_basis_batched(
             beta2=beta2,
             eps=eps,
             max_grad_norm=mgn_eff,
+            frozen_set=frozen_set if frozen_set else None,
         )
 
         # Initialise Adam moment buffers ONCE — they persist across all steps,
@@ -302,6 +362,7 @@ def train_basis_batched(
                     max_iter=1,
                     tol=0.0,
                     record_loss=True,
+                    frozen_indices=frozen_set if frozen_set else None,
                 )
                 loss_history.append(step_trace[-1] if len(step_trace) >= 2 else step_trace[0])
                 global_step += 1

pdft-0.2.2/tests/circuit/test_h_inverse_leg_swap.py

@@ -0,0 +1,142 @@
+"""Regression test for the H-gate inverse leg-swap bug.
+
+The stepped-tensordot rewrite of `compile_circuit` (commit
+`feat(circuit): stepped tensordot for circuits beyond the 52-label pool`)
+forgot to swap the contracted leg of the Hadamard tensor when
+``inverse=True``. Combined with `loss._scalar_loss`'s convention of
+passing ``conj(tensors)`` to the inverse closure, this means the
+inverse path applies ``conj(H)`` instead of the proper adjoint
+``conj(H.T) = H†``.
+
+For the default symmetric Hadamard (``[[1, 1], [1, -1]] / sqrt 2``),
+``conj(H) = H = H†`` so the bug is invisible. For a *trained* (and
+therefore non-symmetric) Hadamard the bug breaks the round-trip:
+``T† T x ≠ x``. The U4 gate handler in the same function does the
+leg-swap correctly already; this test pins the equivalent contract for
+the H gate.
+"""
+
+from __future__ import annotations
+
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+import pdft
+
+
+def _non_symmetric_unitary(seed: int = 0) -> jnp.ndarray:
+    """A 2x2 unitary that is NOT real-symmetric (so conj ≠ adjoint).
+
+    Built as exp(i * theta * X) for a Hermitian X with non-zero off-diagonal
+    imag parts. The result is unitary by construction and demonstrably
+    non-Hermitian and non-symmetric.
+    """
+    rng = np.random.default_rng(seed)
+    # Random Hermitian generator with complex off-diagonal
+    g = rng.normal(size=(2, 2)) + 1j * rng.normal(size=(2, 2))
+    H = (g + g.conj().T) / 2  # Hermitian
+    # exp(i * H) is unitary; for a non-trivial H it is non-symmetric in general
+    eigvals, eigvecs = np.linalg.eigh(H)
+    U = eigvecs @ np.diag(np.exp(1j * eigvals)) @ eigvecs.conj().T
+    return jnp.asarray(U, dtype=jnp.complex128)
+
+
+def test_qft_inverse_round_trip_with_non_symmetric_h():
+    """Round-trip ``T† T x = x`` must hold to machine precision when one or
+    more Hadamard tensors are non-symmetric unitaries (the situation that
+    arises after Riemannian-Adam training perturbs the gates off the
+    initial real-symmetric H).
+    """
+    m, n = 3, 3
+    basis = pdft.QFTBasis(m=m, n=n)
+
+    # Replace the FIRST Hadamard tensor with a non-symmetric unitary.
+    # QFTBasis sorts tensors Hadamards-first, so tensors[0..m+n-1] are H slots.
+    perturbed = list(basis.tensors)
+    U_nonsym = _non_symmetric_unitary(seed=42)
+    perturbed[0] = U_nonsym
+
+    # Sanity: the replacement is unitary but not symmetric.
+    assert np.allclose(
+        np.asarray(U_nonsym) @ np.asarray(U_nonsym).conj().T, np.eye(2), atol=1e-12
+    ), "test setup error: replacement matrix must be unitary"
+    assert not np.allclose(
+        np.asarray(U_nonsym), np.asarray(U_nonsym).T, atol=1e-6
+    ), "test setup error: replacement matrix must be non-symmetric"
+
+    perturbed_basis = pdft.QFTBasis(m=m, n=n, tensors=perturbed)
+
+    rng = np.random.default_rng(0)
+    x = jnp.asarray(
+        rng.standard_normal((2**m, 2**n)) + 1j * rng.standard_normal((2**m, 2**n)),
+        dtype=jnp.complex128,
+    )
+
+    # Forward must preserve norm (operator is unitary).
+    Tx = perturbed_basis.forward_transform(x)
+    np.testing.assert_allclose(
+        float(jnp.sum(jnp.abs(Tx) ** 2)),
+        float(jnp.sum(jnp.abs(x) ** 2)),
+        rtol=1e-10,
+        err_msg="forward operator is not norm-preserving",
+    )
+
+    # Round-trip must recover x to machine precision.
+    rt = perturbed_basis.inverse_transform(Tx)
+    np.testing.assert_allclose(
+        np.asarray(rt),
+        np.asarray(x),
+        atol=1e-10,
+        err_msg="round-trip T† T x != x — H-gate inverse leg-swap bug",
+    )
+
+
+def test_blocked_inverse_round_trip_with_non_symmetric_h():
+    """Same regression for BlockedBasis(QFTBasis, ...) — this is the path
+    actually hit by the trained_blocked_8 cell whose val loss the bug
+    inflated by ~12x.
+    """
+    inner = pdft.QFTBasis(m=3, n=3)
+    perturbed = list(inner.tensors)
+    perturbed[0] = _non_symmetric_unitary(seed=7)
+    perturbed_inner = pdft.QFTBasis(m=3, n=3, tensors=perturbed)
+
+    blocked = pdft.BlockedBasis(inner=perturbed_inner, block_log_m=2, block_log_n=2)
+
+    rng = np.random.default_rng(1)
+    side = 2 ** blocked.m
+    x = jnp.asarray(
+        rng.standard_normal((side, side)) + 1j * rng.standard_normal((side, side)),
+        dtype=jnp.complex128,
+    )
+
+    Tx = blocked.forward_transform(x)
+    np.testing.assert_allclose(
+        float(jnp.sum(jnp.abs(Tx) ** 2)),
+        float(jnp.sum(jnp.abs(x) ** 2)),
+        rtol=1e-10,
+    )
+
+    rt = blocked.inverse_transform(Tx)
+    np.testing.assert_allclose(
+        np.asarray(rt), np.asarray(x), atol=1e-10,
+        err_msg="BlockedBasis round-trip broken — H-gate inverse leg-swap bug",
+    )
+
+
+@pytest.mark.parametrize("seed", [0, 1, 2, 3, 4])
+def test_qft_inverse_round_trip_with_non_symmetric_h_random_seeds(seed: int):
+    """Robustness: hold across multiple random non-symmetric H replacements."""
+    basis = pdft.QFTBasis(m=2, n=2)
+    perturbed = list(basis.tensors)
+    perturbed[0] = _non_symmetric_unitary(seed=seed)
+    perturbed_basis = pdft.QFTBasis(m=2, n=2, tensors=perturbed)
+
+    rng = np.random.default_rng(seed)
+    x = jnp.asarray(
+        rng.standard_normal((4, 4)) + 1j * rng.standard_normal((4, 4)),
+        dtype=jnp.complex128,
+    )
+    rt = perturbed_basis.inverse_transform(perturbed_basis.forward_transform(x))
+    np.testing.assert_allclose(np.asarray(rt), np.asarray(x), atol=1e-10)

pdft-0.2.2/tests/circuit/test_high_qubit_circuit.py

@@ -0,0 +1,62 @@
+"""Regression tests for circuits whose flat-einsum representation would
+exceed the 52-character a-zA-Z label pool.
+
+The pre-fix `build_circuit_einsum` allocates a fresh label per Hadamard
+output and two fresh labels per U4 output. RealRichBasis at inner_m=5,
+inner_n=5 emits 2(m + m²) = 60 fresh labels — over the limit and raising
+ValueError at construction. The stepped-contraction implementation
+applies one gate per `jnp.tensordot`, recycling labels across steps so
+the limit never bites.
+
+These tests pin both:
+  - construction succeeds at inner_m=5
+  - the forward+inverse roundtrip returns the input image to machine
+    precision (a 32×32 image at m=n=5)
+"""
+
+from __future__ import annotations
+
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from pdft import RealRichBasis, RichBasis
+
+
+@pytest.mark.parametrize("BasisCls", [RichBasis, RealRichBasis])
+def test_basis_constructs_at_m5(BasisCls):
+    """Construction must succeed at m=n=5 (10 qubits per block, 60 fresh labels)."""
+    basis = BasisCls(m=5, n=5)
+    assert basis.m == 5
+    assert basis.n == 5
+
+
+@pytest.mark.parametrize("BasisCls", [RichBasis, RealRichBasis])
+def test_basis_forward_inverse_roundtrip_at_m5(BasisCls):
+    """Forward then inverse on the initial (untrained) basis must reproduce the input.
+
+    At step 0 the basis is initialised to a known unitary (QFT-like for RichBasis,
+    Walsh-Hadamard ⊗ identity for RealRichBasis), so U U^{-1} = I exactly.
+    """
+    rng = np.random.default_rng(0)
+    pic = jnp.asarray(rng.normal(size=(32, 32)).astype(np.float64))
+
+    basis = BasisCls(m=5, n=5)
+    transformed = basis.forward_transform(pic)
+    recovered = basis.inverse_transform(transformed)
+
+    np.testing.assert_allclose(np.real(recovered), np.asarray(pic), atol=1e-9)
+
+
+@pytest.mark.parametrize("BasisCls", [RichBasis, RealRichBasis])
+def test_basis_forward_preserves_norm_at_m5(BasisCls):
+    """Unitary forward circuits must preserve the Frobenius norm of the input."""
+    rng = np.random.default_rng(1)
+    pic = jnp.asarray(rng.normal(size=(32, 32)).astype(np.float64))
+
+    basis = BasisCls(m=5, n=5)
+    transformed = basis.forward_transform(pic)
+
+    norm_in = float(jnp.linalg.norm(pic))
+    norm_out = float(jnp.linalg.norm(transformed))
+    np.testing.assert_allclose(norm_out, norm_in, rtol=1e-10)