pdft 0.2.1__tar.gz → 0.2.3__tar.gz

{pdft-0.2.1 → pdft-0.2.3}/.claude/CLAUDE.md

@@ -80,8 +80,8 @@ src/pdft/
 ├── profiling.py            Cross-cutting profiling helpers
 ├── bases/
 │   ├── base.py             AbstractSparseBasis + bases_allclose + 4 basis dataclasses
-│   ├── circuit/            QFT, EntangledQFT, TEBD, MERA (full circuit topologies)
-│   └── block/              Blocked, Rich, RealRich (parameter-efficient)
+│   ├── circuit/            QFT, EntangledQFT, TEBD, MERA, Rich, RealRich + freeze_as_blocked
+│   └── block/              BlockedBasis (Rich/RealRich re-exported from circuit for back-compat)
 ├── circuit/                Einsum builder (builder.py) + JIT closure cache (cache.py)
 ├── optimizers/             core, gd (RiemannianGD + Armijo), adam (RiemannianAdam), loop
 ├── training/               schedules, single (train_basis), batched, adam_step, eval_loop

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pdft
-Version: 0.2.1
+Version: 0.2.3
 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.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "pdft"
-version = "0.2.1"
+version = "0.2.3"
 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.3}/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.3"
 __upstream_ref__ = "nzy1997/ParametricDFT.jl@a201a27e47df2f0f3ab460f83d49b6e5f5d1e9ef"
 
 # Slim public re-export hub: most-used names only. Anything else: import
@@ -49,6 +49,7 @@ __upstream_ref__ = "nzy1997/ParametricDFT.jl@a201a27e47df2f0f3ab460f83d49b6e5f5d
 from .bases import (  # noqa: E402
     AbstractSparseBasis,
     BlockedBasis,
+    DCT4Basis,
     EntangledQFTBasis,
     MERABasis,
     QFTBasis,
@@ -56,6 +57,7 @@ from .bases import (  # noqa: E402
     RichBasis,
     TEBDBasis,
     bases_allclose,
+    freeze_as_blocked,
 )
 from .loss import AbstractLoss, L1Norm, MSELoss, loss_function  # noqa: E402
 from .optimizers import RiemannianAdam, RiemannianGD, optimize  # noqa: E402
@@ -65,6 +67,7 @@ __all__ = [
     "AbstractLoss",
     "AbstractSparseBasis",
     "BlockedBasis",
+    "DCT4Basis",
     "EntangledQFTBasis",
     "L1Norm",
     "MERABasis",
@@ -79,6 +82,7 @@ __all__ = [
     "__upstream_ref__",
     "__version__",
     "bases_allclose",
+    "freeze_as_blocked",
     "loss_function",
     "optimize",
     "train_basis",

{pdft-0.2.1 → pdft-0.2.3}/src/pdft/bases/__init__.py

@@ -3,26 +3,39 @@
 Two families:
 - bases.circuit  — full circuit topologies (QFT, EntangledQFT, TEBD, MERA),
                    comparable to FFT/DCT.
-- bases.block    — parameter-efficient block-structured bases (Blocked,
-                   Rich, RealRich) over arbitrary block partitions.
+- bases.block    — BlockedBasis over arbitrary block partitions (Rich/RealRich
+                   are re-exported here for back-compat; they live in circuit).
 
 The abstract base class and bases_allclose helper live in bases.base.
 """
 
 from .base import (
     AbstractSparseBasis,
+    DCT4Basis,
     EntangledQFTBasis,
     MERABasis,
     QFTBasis,
     TEBDBasis,
     bases_allclose,
 )
-from .block import BlockedBasis, RealRichBasis, RichBasis, fit_to_dct
-from .circuit import entangled_qft_code, ft_mat, ift_mat, mera_code, qft_code, tebd_code
+from .block import BlockedBasis
+from .circuit import (
+    RealRichBasis,
+    RichBasis,
+    entangled_qft_code,
+    fit_to_dct,
+    freeze_as_blocked,
+    ft_mat,
+    ift_mat,
+    mera_code,
+    qft_code,
+    tebd_code,
+)
 
 __all__ = [
     "AbstractSparseBasis",
     "BlockedBasis",
+    "DCT4Basis",
     "EntangledQFTBasis",
     "MERABasis",
     "QFTBasis",
@@ -32,6 +45,7 @@ __all__ = [
     "bases_allclose",
     "entangled_qft_code",
     "fit_to_dct",
+    "freeze_as_blocked",
     "ft_mat",
     "ift_mat",
     "mera_code",

{pdft-0.2.1 → pdft-0.2.3}/src/pdft/bases/base.py

@@ -291,11 +291,15 @@ class TEBDBasis:
         code: object | None = None,
         inv_code: object | None = None,
         seed: int | None = None,
+        parametrization: str = "cp",
     ):
         """When `seed` is provided AND `phases` is None, draws phases from
         `np.random.default_rng(seed).normal(0, 0.1, n_row_gates + n_col_gates)`
         — Julia's `randn(n_gates) * 0.1` convention from
         `ParametricDFT.jl/src/training.jl::_init_circuit`.
+
+        `parametrization` is "cp" (diagonal ring gates on U(1)^4) or "u4"
+        (dense two-qubit ring gates on U(4), the canonical TEBD gate).
         """
         import numpy as np
 
@@ -307,8 +311,12 @@ class TEBDBasis:
         self.n = n
         if phases is None and seed is not None:
             phases = list(np.random.default_rng(seed).normal(0.0, 0.1, m + n))
-        _code, init_tensors, self.n_row_gates, self.n_col_gates = tebd_code(m, n, phases=phases)
-        _inv_code, _, _, _ = tebd_code(m, n, phases=phases, inverse=True)
+        _code, init_tensors, self.n_row_gates, self.n_col_gates = tebd_code(
+            m, n, phases=phases, parametrization=parametrization
+        )
+        _inv_code, _, _, _ = tebd_code(
+            m, n, phases=phases, inverse=True, parametrization=parametrization
+        )
         self.tensors = list(tensors) if tensors is not None else init_tensors
         self.code = code if code is not None else _code
         self.inv_code = inv_code if inv_code is not None else _inv_code
@@ -394,11 +402,15 @@ class MERABasis:
         code: object | None = None,
         inv_code: object | None = None,
         seed: int | None = None,
+        parametrization: str = "cp",
     ):
         """When `seed` is provided AND `phases` is None, draws phases from
         `np.random.default_rng(seed).normal(0, 0.1, n_gates)` — Julia's
         `randn(n_gates) * 0.1` convention from
         `ParametricDFT.jl/src/training.jl::_init_circuit`.
+
+        `parametrization` is "cp" (diagonal disentanglers/isometries on
+        U(1)^4) or "u4" (dense two-qubit gates on U(4), the canonical form).
         """
         import numpy as np
 
@@ -412,8 +424,12 @@ class MERABasis:
             n_row = _n_mera_gates(m) if m >= 2 else 0
             n_col = _n_mera_gates(n) if n >= 2 else 0
             phases = list(np.random.default_rng(seed).normal(0.0, 0.1, n_row + n_col))
-        _code, init_tensors, self.n_row_gates, self.n_col_gates = mera_code(m, n, phases=phases)
-        _inv_code, _, _, _ = mera_code(m, n, phases=phases, inverse=True)
+        _code, init_tensors, self.n_row_gates, self.n_col_gates = mera_code(
+            m, n, phases=phases, parametrization=parametrization
+        )
+        _inv_code, _, _, _ = mera_code(
+            m, n, phases=phases, inverse=True, parametrization=parametrization
+        )
         self.tensors = list(tensors) if tensors is not None else init_tensors
         self.code = code if code is not None else _code
         self.inv_code = inv_code if inv_code is not None else _inv_code
@@ -468,3 +484,107 @@ def _mera_unflatten(aux, leaves) -> MERABasis:
 
 
 tree_util.register_pytree_node(MERABasis, _mera_flatten, _mera_unflatten)
+
+
+# ---------------------------------------------------------------------------
+# DCT4Basis
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class DCT4Basis:
+    """DCT-IV tensor-network basis: the real-orthogonal, ancilla-free analogue
+    of :class:`QFTBasis`.
+
+    Mirrors QFTBasis (one tensor list; `inverse_transform` applies
+    `conj(tensors)` through `inv_code`). The gates are emitted by
+    :func:`pdft.bases.circuit.dct4.dct4_code`; at initialization the forward
+    operator is the bit-reversed orthonormal DCT-IV per dimension, and — since
+    DCT-IV is self-inverse — the basis reconstructs exactly. ``tensors`` holds
+    every gate — the affine ``R_y`` rotation layer, the branch Hadamards, the
+    mirror-``Q`` CNOT permutations and the ``Delta`` sign — each a learnable
+    leaf on its auto-selected Riemannian manifold (O(2) / O(4) / phase). The
+    gate tensors are real-valued (stored complex128, zero imaginary), so the
+    unitary manifold trains the real-orthogonal subset under a real objective:
+    exact DCT-IV at init, then relaxed, just as QFT relaxes within U.
+
+    Pytree contract:
+        leaves   = tensors
+        aux data = (m, n, len(tensors), code, inv_code)
+    """
+
+    m: int
+    n: int
+    tensors: list[Array]
+    code: object = field(compare=False, repr=False)
+    inv_code: object = field(compare=False, repr=False)
+
+    def __init__(
+        self,
+        m: int,
+        n: int,
+        tensors: Sequence[Array] | None = None,
+        code: object | None = None,
+        inv_code: object | None = None,
+        parametrization: str = "o4",
+    ):
+        if m < 1 or n < 1:
+            raise ValueError(f"m and n must be >= 1, got m={m}, n={n}")
+        from .circuit.dct4 import dct4_code
+
+        self.m = m
+        self.n = n
+        _code, init_tensors = dct4_code(m, n, parametrization=parametrization)
+        _inv_code, _ = dct4_code(m, n, inverse=True, parametrization=parametrization)
+        self.tensors = list(tensors) if tensors is not None else init_tensors
+        self.code = code if code is not None else _code
+        self.inv_code = inv_code if inv_code is not None else _inv_code
+
+    @property
+    def inv_tensors(self) -> list[Array]:
+        return self.tensors
+
+    @property
+    def image_size(self) -> tuple[int, int]:
+        return (2**self.m, 2**self.n)
+
+    @property
+    def num_parameters(self) -> int:
+        return sum(int(t.size) for t in self.tensors)
+
+    def forward_transform(self, pic: Array) -> Array:
+        from .circuit.dct4 import dct4_ft_mat
+
+        return dct4_ft_mat(self.tensors, self.code, self.m, self.n, pic)
+
+    def inverse_transform(self, pic: Array) -> Array:
+        from .circuit.dct4 import dct4_ift_mat
+
+        return dct4_ift_mat(
+            [jnp.conj(t) for t in self.tensors],
+            self.inv_code,
+            self.m,
+            self.n,
+            pic,
+        )
+
+
+def _dct4basis_flatten(b: DCT4Basis):
+    leaves = tuple(b.tensors)
+    aux = (b.m, b.n, len(b.tensors), b.code, b.inv_code)
+    return leaves, aux
+
+
+def _dct4basis_unflatten(aux, leaves) -> DCT4Basis:
+    m, n, n_fwd, code, inv_code = aux
+    assert len(leaves) == n_fwd
+    out = DCT4Basis.__new__(DCT4Basis)
+    out.m = m
+    out.n = n
+    out.tensors = list(leaves)
+    out.code = code
+    out.inv_code = inv_code
+    return out
+
+
+tree_util.register_pytree_node(DCT4Basis, _dct4basis_flatten, _dct4basis_unflatten)

pdft-0.2.3/src/pdft/bases/block/__init__.py

@@ -0,0 +1,15 @@
+"""Block-structured bases: BlockedBasis.
+
+RichBasis/RealRichBasis/fit_to_dct now live in bases.circuit (they are full
+circuits); re-exported here for back-compat.
+"""
+
+from ..circuit import RealRichBasis, RichBasis, fit_to_dct
+from .block import BlockedBasis
+
+__all__ = [
+    "BlockedBasis",
+    "RealRichBasis",
+    "RichBasis",
+    "fit_to_dct",
+]

pdft-0.2.3/src/pdft/bases/circuit/__init__.py

@@ -0,0 +1,30 @@
+"""Circuit-topology bases: QFT, DCT-IV, EntangledQFT, TEBD, MERA, Rich, RealRich.
+
+Rich/RealRich are H + U(4) circuits on the QFT topology — full-image
+transforms usable standalone or as a BlockedBasis inner.
+"""
+
+from .qft import ft_mat, ift_mat, qft_code
+from .dct4 import dct4_code, dct4_ft_mat, dct4_ift_mat
+from .entangled_qft import entangled_qft_code
+from .mera import mera_code
+from .tebd import tebd_code
+from .rich import RichBasis, fit_to_dct
+from .real_rich import RealRichBasis
+from .freeze import freeze_as_blocked
+
+__all__ = [
+    "RealRichBasis",
+    "RichBasis",
+    "dct4_code",
+    "dct4_ft_mat",
+    "dct4_ift_mat",
+    "entangled_qft_code",
+    "fit_to_dct",
+    "freeze_as_blocked",
+    "ft_mat",
+    "ift_mat",
+    "mera_code",
+    "qft_code",
+    "tebd_code",
+]

pdft-0.2.3/src/pdft/bases/circuit/dct4.py

@@ -0,0 +1,189 @@
+"""DCT-IV circuit as a hand-rolled tensor network.
+
+Real-orthogonal analogue of the QFT (cf. :mod:`pdft.bases.circuit.qft`). The
+DCT-IV is the self-dual "bulk" cosine transform: real, orthogonal, and its
+own inverse. Unlike the DCT-II it is **ancilla-free** and uses only 1- and
+2-qubit gates, so it drops straight into the H + 2-qubit-gate tensor-network
+builder with no multi-controlled scaffolding.
+
+The 1D recursion is *uniform* — both halves run the same child, so DCT-IV
+unrolls into n qubit-local stages with no multi-controls (the real-orthogonal
+analogue of the QFT butterfly). For N = 2M, with top qubit b splitting the
+register (cf. ``all_real_dct_one_ancilla.tex`` §3.2):
+
+    C^IV_{2M} = P_{2M} (H (x) I_M) (I_M (+) D_M) (I_2 (x) C^IV_M)
+                (I_M (+) R_M) T_{2M} Q_{2M}
+
+Per level k on the within-block qubits (b = top qubit), in emission order
+(= the operator product above read right-to-left):
+
+    * Q  pair layout : CX(control=b, target=q) for each q below b
+    * T  rotation    : R_y(pi / 2N) on b, then CR_y(pi 2^p / N) on b controlled
+                       by each lower bit — the affine phase gradient (one base
+                       rotation + one controlled rotation per bit, never a 2^m
+                       multiplexer)
+    * R  reverse     : the mirror-Q permutation repeated (odd-branch reversal)
+    * recurse        : C^IV_M on the lower register, uncontrolled (I_2 (x) C^IV_M)
+    * D  sign Delta  : CZ(control=b, target=next qubit)
+    * H  merge       : H on b
+
+The interleave P_{2M} is deferred: per-level interleaves compose into a single
+bit reversal, realised as the bit-reversed output convention (matching
+:func:`qft_code`), so no explicit P gate is emitted.
+
+At init the forward operator equals the bit-reversed orthonormal DCT-IV per
+dimension. Every gate is a learnable leaf on its auto-selected Riemannian
+manifold — R_y and branch-H on O(2), CR_y and mirror-Q on O(4), the Delta
+sign on the phase manifold. Like QFT/RealRich, real init plus a real objective
+keep the operator real-orthogonal as it relaxes from the exact transform (cf.
+``all_real_dct_zero_ancilla.tex``).
+
+2D DCT-IV = (m-qubit DCT-IV on row qubits) tensor (n-qubit DCT-IV on col
+qubits); no entanglement between blocks.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from ...circuit.builder import (
+    HADAMARD,
+    Gate,
+    apply_circuit,
+    compile_circuit,
+    controlled_phase_diag,
+)
+
+Array = jax.Array
+
+
+__all__ = [
+    "dct4_code",
+    "dct4_ft_mat",
+    "dct4_ift_mat",
+    "_dct4_gates_1d",
+    "_cry",
+]
+
+
+def _ry(theta: float) -> Array:
+    """Single-qubit real R_y(theta) as a (2, 2) complex tensor `T[out, in]`."""
+    c, s = np.cos(theta / 2.0), np.sin(theta / 2.0)
+    return jnp.asarray(np.array([[c, -s], [s, c]]), dtype=jnp.complex128)
+
+
+def _cnot_u4() -> Array:
+    """CNOT as a (2, 2, 2, 2) tensor `T[out_ctrl, out_tgt, in_ctrl, in_tgt]`."""
+    M = np.zeros((2, 2, 2, 2))
+    for ic in (0, 1):
+        for it in (0, 1):
+            M[ic, it ^ ic, ic, it] = 1.0
+    return jnp.asarray(M, dtype=jnp.complex128)
+
+
+def _cry_u4(theta: float) -> Array:
+    """Controlled-R_y(theta) as a (2, 2, 2, 2) tensor (control fires on 1)."""
+    c, s = np.cos(theta / 2.0), np.sin(theta / 2.0)
+    R = np.array([[c, -s], [s, c]])
+    M = np.zeros((2, 2, 2, 2))
+    for it in (0, 1):
+        M[0, it, 0, it] = 1.0  # control = 0 -> identity on target
+    for ot in (0, 1):
+        for it in (0, 1):
+            M[1, ot, 1, it] = R[ot, it]  # control = 1 -> R_y on target
+    return jnp.asarray(M, dtype=jnp.complex128)
+
+
+def _cry(theta: float) -> Array:
+    """Controlled-R_y trainable leaf: just the (2, 2) R_y block.
+
+    Applied with control structure by the builder's ``CRY`` kind (identity on
+    control = 0, this block on control = 1), so the leaf trains on O(2) — the
+    structured, dense-O(4)-free parametrization of the twiddle.
+    """
+    return _ry(theta)
+
+
+def _dct4_gates_1d(n_qubits: int, offset: int, parametrization: str = "o4") -> list[Gate]:
+    """Emit the 1D DCT-IV gate sequence on qubits (offset+1, ..., offset+n_qubits).
+
+    Within-block qubit ``q`` (0 = top/most significant) maps to builder qubit
+    ``offset + n_qubits - q`` (Yao little-endian), which reproduces the
+    bit-reversed orthonormal DCT-IV exactly.
+    """
+
+    def Q(q: int) -> int:
+        return offset + n_qubits - q
+
+    gates: list[Gate] = []
+
+    def level(k: int) -> None:
+        if k == n_qubits:
+            return
+        size = 2 ** (n_qubits - k)
+        b = k
+        lower = list(range(k + 1, n_qubits))
+        # Q: pair-layout permutation — controlled bit-complement of the lower
+        # register (exact CNOT at init; a learnable leaf like every gate)
+        for q in lower:
+            gates.append(Gate(kind="U4", qubits=(Q(b), Q(q)), tensor=_cnot_u4(), phase=0.0))
+        # T: affine R_y phase-gradient (one base R_y on b, one CR_y per lower bit)
+        gates.append(Gate(kind="H", qubits=(Q(b),), tensor=_ry(np.pi / (2 * size)), phase=0.0))
+        for p in range(n_qubits - 1 - k):
+            theta = np.pi * (2**p) / size
+            if parametrization == "controlled":
+                gates.append(Gate(kind="CRY", qubits=(Q(n_qubits - 1 - p), Q(b)),
+                                  tensor=_cry(theta), phase=0.0))
+            else:
+                gates.append(Gate(kind="U4", qubits=(Q(n_qubits - 1 - p), Q(b)),
+                                  tensor=_cry_u4(theta), phase=0.0))
+        # R: odd-branch reversal — the mirror-Q permutation repeated
+        for q in lower:
+            gates.append(Gate(kind="U4", qubits=(Q(b), Q(q)), tensor=_cnot_u4(), phase=0.0))
+        # recurse on the lower register (uncontrolled)
+        level(k + 1)
+        # D: Delta sign  +  H: branch Hadamard merge
+        if lower:
+            gates.append(
+                Gate(kind="CP", qubits=(Q(b), Q(k + 1)), tensor=controlled_phase_diag(float(np.pi)), phase=float(np.pi))
+            )
+        gates.append(Gate(kind="H", qubits=(Q(b),), tensor=HADAMARD, phase=0.0))
+
+    level(0)
+    return gates
+
+
+def dct4_code(
+    m: int, n: int, *, inverse: bool = False, parametrization: str = "o4"
+) -> tuple[Callable[..., Array], list[Array]]:
+    """Return `(einsum_fn, initial_tensors)` for 2D DCT-IV on (2^m, 2^n) images.
+
+    ``parametrization`` selects how the affine twiddle is stored: ``"o4"``
+    (default) emits a dense ``(2, 2, 2, 2)`` controlled-R_y trained on O(4);
+    ``"controlled"`` emits a single-angle ``CRY`` gate whose trainable leaf is
+    a ``(2, 2)`` block on O(2) (the mirror-Q CNOTs stay dense U4).
+    """
+    if m < 1 or n < 1:
+        raise ValueError(f"m and n must be >= 1, got m={m}, n={n}")
+    if parametrization not in ("o4", "controlled"):
+        raise ValueError(
+            f"parametrization must be 'o4' or 'controlled', got {parametrization!r}"
+        )
+    gates = _dct4_gates_1d(m, offset=0, parametrization=parametrization) + _dct4_gates_1d(
+        n, offset=m, parametrization=parametrization
+    )
+    return compile_circuit(gates, m, n, inverse=inverse)
+
+
+def dct4_ft_mat(tensors: list[Array], code: Callable, m: int, n: int, pic: Array) -> Array:
+    """Apply 2D DCT-IV circuit to a (2^m, 2^n) image."""
+    return apply_circuit(tensors, code, m, n, pic)
+
+
+def dct4_ift_mat(tensors: list[Array], code: Callable, m: int, n: int, pic: Array) -> Array:
+    """Apply 2D inverse DCT-IV circuit. Caller must have conjugated the tensors."""
+    return apply_circuit(tensors, code, m, n, pic)

pdft-0.2.3/src/pdft/bases/circuit/freeze.py

@@ -0,0 +1,116 @@
+"""freeze_as_blocked: reduce a full circuit to its blocked equivalent.
+
+Resets the outer (block-index) gates of a QFT/Rich/RealRich circuit to
+identity and returns the frozen tensor indices, so that
+
+    train_basis_batched(frozen_basis, frozen_indices=frozen, ...)
+
+reproduces BlockedBasis(inner, block_log_m, block_log_n) training dynamics.
+See docs/superpowers/specs/2026-05-24-circuit-rich-frozen-blocked-design.md.
+
+Gradient-norm clipping note: frozen slots have their Euclidean gradient zeroed
+before manifold projection, so their Riemannian gradient is zero and they
+contribute zero to the global grad-norm. The clip factor therefore matches
+BlockedBasis exactly — parity holds with max_grad_norm set (see
+test_frozen_qft_outer_gates_matches_blocked_qft_training at max_grad_norm=1.0).
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+
+from ...circuit.builder import controlled_phase_diag, sorted_gate_program
+
+Array = jax.Array
+
+__all__ = ["freeze_as_blocked"]
+
+
+def _identity_for_kind(kind: str) -> Array:
+    """Identity-acting tensor for a gate kind, as complex128."""
+    if kind == "H":
+        return jnp.eye(2, dtype=jnp.complex128)
+    if kind == "U4":
+        return jnp.eye(4, dtype=jnp.complex128).reshape(2, 2, 2, 2)
+    if kind == "CP":
+        return controlled_phase_diag(0.0)  # diag(1,1,1,1) in compact 2x2 form
+    raise AssertionError(f"unknown gate kind: {kind}")
+
+
+def freeze_as_blocked(basis: Any, block_log_m: int, block_log_n: int) -> tuple[Any, list[int]]:
+    """Return ``(basis_copy_with_identity_outer_gates, frozen_indices)``.
+
+    Training the returned basis with ``frozen_indices`` reproduces
+    ``BlockedBasis(inner, block_log_m, block_log_n)`` training dynamics, where
+    ``inner`` is the same basis type at ``(m - block_log_m, n - block_log_n)``.
+    The outer (block-index) gates are reset to identity; inner gates keep the
+    input basis's tensor values.
+
+    Supports ``QFTBasis``, ``RichBasis``, ``RealRichBasis``.
+
+    ``block_log_m == block_log_n == 0`` is a valid no-op: no qubits are
+    block-index qubits, so nothing is frozen and ``frozen_indices`` is empty
+    (the returned basis is a tensor-copy of the input).
+    """
+    # Lazy imports avoid a real bases.base <-> bases.circuit import cycle.
+    # base.py runs `from .circuit.qft import qft_code` at module load; importing
+    # that submodule first executes bases/circuit/__init__.py, which imports
+    # THIS module. A module-top `from ..base import QFTBasis` would then touch
+    # base.py while it is still partially initialised (before QFTBasis exists)
+    # and raise ImportError. Verified: hoisting these to the top breaks
+    # `import pdft`. Keep them inside the function.
+    from ..base import QFTBasis
+    from .qft import _qft_gates_1d
+    from .real_rich import RealRichBasis, _real_rich_qft_gates_1d
+    from .rich import RichBasis, _rich_qft_gates_1d
+
+    gate_builders = {
+        QFTBasis: _qft_gates_1d,
+        RichBasis: _rich_qft_gates_1d,
+        RealRichBasis: _real_rich_qft_gates_1d,
+    }
+    btype = type(basis)
+    if btype not in gate_builders:
+        raise TypeError(
+            "freeze_as_blocked supports QFTBasis, RichBasis, RealRichBasis; "
+            f"got {btype.__name__}"
+        )
+    if block_log_m < 0 or block_log_n < 0:
+        raise ValueError(
+            "block_log_m and block_log_n must be >= 0; got "
+            f"block_log_m={block_log_m}, block_log_n={block_log_n}"
+        )
+    m, n = basis.m, basis.n
+    if m - block_log_m < 1 or n - block_log_n < 1:
+        raise ValueError(
+            f"block partition leaves no inner qubits: m={m}, n={n}, "
+            f"block_log_m={block_log_m}, block_log_n={block_log_n} "
+            "(need m - block_log_m >= 1 and n - block_log_n >= 1)"
+        )
+
+    # Block-index qubits: higher-numbered qubits per dim (Yao little-endian),
+    # matching BlockedBasis.
+    block_qubits = set(range(m - block_log_m + 1, m + 1)) | set(
+        range(m + n - block_log_n + 1, m + n + 1)
+    )
+
+    builder = gate_builders[btype]
+    gates = builder(m, offset=0) + builder(n, offset=m)
+    program = sorted_gate_program(gates)
+    if len(program) != len(basis.tensors):
+        raise AssertionError(
+            f"gate program length {len(program)} != tensor count {len(basis.tensors)}"
+        )
+
+    new_tensors = [jnp.array(t, copy=True) for t in basis.tensors]
+    frozen_indices: list[int] = []
+    for i, (kind, qubits) in enumerate(program):
+        if set(qubits) & block_qubits:
+            new_tensors[i] = _identity_for_kind(kind).astype(new_tensors[i].dtype)
+            frozen_indices.append(i)
+
+    new_basis = btype(m=m, n=n, tensors=new_tensors, code=basis.code, inv_code=basis.inv_code)
+    return new_basis, frozen_indices