From f1827f243a32024163ae03e0ce638839f35bf7fb Mon Sep 17 00:00:00 2001
From: Panadestein <rpana92@gmail.com>
Date: Tue, 1 Sep 2026 16:54:25 +0000
Subject: [PATCH 1/4] =?UTF-8?q?refactor!:=20=F0=9F=94=A5=20drop=20quimb=20?=
 =?UTF-8?q?as=20a=20runtime=20dependency?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

quimb contributed nothing to the SRC algorithm itself: every hot path is
opt_einsum over plain numpy/cupy arrays. It was used only as a container
type (isinstance dispatch, attribute access, re-wrapping results) and for
the sub-3-site SVD fallback, yet it pulled cotengra, cytoolz, psutil,
scipy and tqdm into every install.

BREAKING CHANGE: `apply` and `compress` now take and return plain lists of
per-site arrays instead of quimb MPS/MPO objects. The array layout is
unchanged (default quimb index ordering), so callers can round-trip with
`qtn.MatrixProductOperator(result)` / `qtn.MatrixProductState(result)`.

- infer MPS vs MPO from the rank of the first site tensor
- replace the quimb sub-3-site fallback with an exact two-site SVD
- move quimb to the `test` dependency group, where it is still used to
  build reference networks and measure distances
---
 CONTRIBUTING.md                              |   2 +-
 README.md                                    |  14 +-
 SECURITY.md                                  |   2 +-
 benches/primitives/leonardo/bench_mpo_mpo.py |   6 +-
 docs/index.md                                |  14 +-
 pyproject.toml                               |  10 +-
 src/src_method/_tensor_train.py              | 164 +++++++++++++++++
 src/src_method/apply.py                      | 101 ++++++-----
 src/src_method/compress.py                   |  73 ++++----
 tests/test_gpu_backend.py                    |  68 ++++++-
 tests/test_package.py                        | 177 ++++++++++++++-----
 uv.lock                                      |  50 +++---
 12 files changed, 507 insertions(+), 174 deletions(-)
 create mode 100644 src/src_method/_tensor_train.py

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 2192fe7..52a8b9b 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -14,7 +14,7 @@ comment on that thread and applies to all your future contributions.
 
 Open an issue with:
 
-- the versions of `src_method`, `numpy`, `quimb` and (if relevant) `cupy`,
+- the versions of `src_method`, `numpy` and (if relevant) `cupy`,
 - a minimal reproducer, ideally with a fixed `seed=`,
 - the observed and expected behaviour.
 
diff --git a/README.md b/README.md
index 69fc66c..d146457 100644
--- a/README.md
+++ b/README.md
@@ -15,7 +15,7 @@ The following primitives are supported:
 3. MPO randomized compression.
 4. MPS randomized compression.
 
-The package is designed to work with [Quimb](https://quimb.readthedocs.io/en/latest/autoapi/quimb/tensor/index.html) tensor network objects, as such, the API loosely follows its naming conventions:
+`src_method` has no tensor-network framework dependency: it takes and returns plain lists of per-site NumPy arrays, one array per site.
 
 ```python
 from src_method import apply, compress
@@ -25,11 +25,19 @@ The `apply` function covers cases 1 and 2 above, while the `compress` function c
 manage the assignment of the returned objects, possibly overwriting the input variables.
 See the [reference documentation](algorithmiq.github.io/src_method/) for details, and the [tests](tests/) or [benchmarks](benches/) folders for usage examples.
 
-**NOTE**: the current implementation targets tensor networks with 3 or more sites. For smaller networks, the corresponding Quimb primitive with randomized SVD is dispatched, with a warning.
+Whether a train is an MPS or an MPO is inferred from the rank of its first site tensor, so no wrapper type is needed.
+
+**NOTE**: the current implementation targets tensor networks with 3 or more sites. For smaller networks, an exact SVD-based fallback is dispatched, with a warning.
 
 ### Tensor Indexing Conventions
 
-This library follows the default `quimb` tensor indexing conventions.
+The array layout follows the default `quimb` tensor indexing conventions, so results round-trip through [Quimb](https://quimb.readthedocs.io/en/latest/autoapi/quimb/tensor/index.html) without any permutation:
+
+```python
+import quimb.tensor as qtn
+
+result = qtn.MatrixProductOperator(apply(H1.arrays, H2.arrays, chi_out=64))
+```
 
 - **MPO Tensors:** Bulk tensors have index order `('l', 'r', 'u', 'd')`.
   Boundary tensors (at the edges) are rank-3, dropping the outer `'l'` or `'r'` index.
diff --git a/SECURITY.md b/SECURITY.md
index 913ff70..9c665a8 100644
--- a/SECURITY.md
+++ b/SECURITY.md
@@ -23,4 +23,4 @@ arrays supplied by the caller and does not parse untrusted input formats,
 open network connections, or execute user-supplied code. The most likely
 security-relevant issues are therefore memory-safety problems surfaced through
 the optional CuPy backend, or dependency vulnerabilities. Reports about the
-behaviour of `numpy`, `quimb` or `cupy` themselves should go to those projects.
+behaviour of `numpy` or `cupy` themselves should go to those projects.
diff --git a/benches/primitives/leonardo/bench_mpo_mpo.py b/benches/primitives/leonardo/bench_mpo_mpo.py
index ebcae2c..2247cfe 100644
--- a/benches/primitives/leonardo/bench_mpo_mpo.py
+++ b/benches/primitives/leonardo/bench_mpo_mpo.py
@@ -78,7 +78,11 @@ def main(
     if run == "src":
         logger.info("Computing SRC's MPO-MPO contraction (with compression)...")
         tms = perf_counter_ns()
-        H_src = apply(H1, H2, chi_out=chi_out, dtype=array_type)
+        # src_method takes and returns plain lists of site arrays; quimb is only
+        # used here to build the inputs and to measure the distance.
+        H_src = qtn.MatrixProductOperator(
+            apply(H1.arrays, H2.arrays, chi_out=chi_out, dtype=array_type)
+        )
         tms = perf_counter_ns() - tms
         logger.info(" SRC's contraction-compression took %s s", tms * 1e-9)
         if compare == "yes":
diff --git a/docs/index.md b/docs/index.md
index 9fa8bc5..f88dfaa 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -13,7 +13,7 @@ The following primitives are supported:
 3. MPO randomized compression.
 4. MPS randomized compression.
 
-The package is designed to work with [Quimb](https://quimb.readthedocs.io/en/latest/autoapi/quimb/tensor/index.html) tensor network objects, as such, the API loosely follows its naming conventions:
+`src_method` has no tensor-network framework dependency: it takes and returns plain lists of per-site NumPy arrays, one array per site.
 
 ```python
 from src_method import apply, compress
@@ -23,11 +23,19 @@ The `apply` function covers cases 1 and 2 above, while the `compress` function c
 manage the assignment of the returned objects, possibly overwriting the input variables.
 See the [reference documentation](algorithmiq.github.io/src_method/) for details, and the [tests](../tests/) or [benchmarks](../benches/) folders for usage examples.
 
-**NOTE**: the current implementation targets tensor networks with 3 or more sites. For smaller networks, the corresponding Quimb primitive with randomized SVD is dispatched, with a warning.
+Whether a train is an MPS or an MPO is inferred from the rank of its first site tensor, so no wrapper type is needed.
+
+**NOTE**: the current implementation targets tensor networks with 3 or more sites. For smaller networks, an exact SVD-based fallback is dispatched, with a warning.
 
 ### Tensor Indexing Conventions
 
-This library follows the default `quimb` tensor indexing conventions.
+The array layout follows the default `quimb` tensor indexing conventions, so results round-trip through [Quimb](https://quimb.readthedocs.io/en/latest/autoapi/quimb/tensor/index.html) without any permutation:
+
+```python
+import quimb.tensor as qtn
+
+result = qtn.MatrixProductOperator(apply(H1.arrays, H2.arrays, chi_out=64))
+```
 
 - **MPO Tensors:** Bulk tensors have index order `('l', 'r', 'u', 'd')`.
   Boundary tensors (at the edges) are rank-3, dropping the outer `'l'` or `'r'` index.
diff --git a/pyproject.toml b/pyproject.toml
index bbd881c..2a00f1e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,7 +39,6 @@ dependencies = [
   "numpy>=2.4.4,<3",
   "opt_einsum>=3.4.0",
   "structlog>=25.5.0",
-  "quimb>=1.14.0,<2",
 ]
 
 
@@ -67,7 +66,14 @@ Discussions = "https://github.com/Algorithmiq/src-method/discussions"
 
 
 [dependency-groups]
-test = ["pytest-cov>=4.0", "pytest>=9.1.0,<10.0", "pytest-benchmark>=5.2.3"]
+# quimb is *not* a runtime dependency: it is only used to build reference
+# tensor networks and to measure distances in the tests and benchmarks.
+test = [
+  "pytest-cov>=4.0",
+  "pytest>=9.1.0,<10.0",
+  "pytest-benchmark>=5.2.3",
+  "quimb>=1.14.0,<2",
+]
 dev = ["scalene", "pre-commit>=3", "ruff==0.16.5", { include-group = "test" }]
 interactive = [
   "ipykernel>=7.2.0",
diff --git a/src/src_method/_tensor_train.py b/src/src_method/_tensor_train.py
new file mode 100644
index 0000000..37510da
--- /dev/null
+++ b/src/src_method/_tensor_train.py
@@ -0,0 +1,164 @@
+"""Array-list tensor-train conventions and exact small-system primitives.
+
+Tensor trains are plain lists of arrays, one per site.  The index ordering
+matches the default `quimb` layout, so a result can be handed straight to
+``qtn.MatrixProductState(arrays)`` / ``qtn.MatrixProductOperator(arrays)``
+without any permutation:
+
+* MPS: ``(bond_r, phys)``, ``(bond_l, bond_r, phys)``, ..., ``(bond_l, phys)``
+* MPO: ``(bond_r, up, down)``, ``(bond_l, bond_r, up, down)``, ...,
+  ``(bond_l, up, down)``
+
+The SRC sweep needs at least three sites, so two-site trains are handled here
+instead.  At that size the whole network fits in a single dense matrix, and one
+exact SVD is both cheaper and more accurate than a randomized sketch.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+import numpy as np
+from opt_einsum import contract
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from numpy.typing import NDArray
+
+# Minimum number of sites for which the randomized SRC sweep is defined.
+MIN_SRC_SITES = 3
+
+# Rank of a boundary (first / last) site tensor, which identifies the train type.
+_MPS_BOUNDARY_NDIM = 2
+_MPO_BOUNDARY_NDIM = 3
+
+TrainKind = Literal["mps", "mpo"]
+
+__all__ = [
+    "MIN_SRC_SITES",
+    "TrainKind",
+    "exact_apply",
+    "exact_compress",
+    "infer_kind",
+]
+
+
+def infer_kind(arrays: Sequence[NDArray]) -> TrainKind | None:
+    """Classify a tensor train from the rank of its first site tensor.
+
+    A boundary site carries one bond index plus either a single physical
+    index (MPS) or an upper/lower pair (MPO), so the rank is unambiguous.
+
+    Args:
+        arrays: The site tensors of the train.
+
+    Returns:
+        ``"mps"``, ``"mpo"``, or ``None`` if the layout is unrecognised.
+    """
+    if len(arrays) == 0:
+        return None
+    ndim = np.ndim(arrays[0])
+    if ndim == _MPS_BOUNDARY_NDIM:
+        return "mps"
+    if ndim == _MPO_BOUNDARY_NDIM:
+        return "mpo"
+    return None
+
+
+def exact_compress(
+    arrays: Sequence[NDArray], chi_out: int, kind: TrainKind
+) -> list[NDArray]:
+    """Compress a two-site train exactly via a single truncated SVD.
+
+    Args:
+        arrays: The two site tensors of the train.
+        chi_out: The maximum bond dimension to keep.
+        kind: Whether the train is an ``"mps"`` or an ``"mpo"``.
+
+    Returns:
+        The compressed train, in right-canonical form.
+
+    Raises:
+        ValueError: If the train does not have exactly two sites.
+    """
+    _check_pair(arrays)
+    if kind == "mps":
+        # (b, p0) x (b, p1) -> (p0, p1)
+        theta = contract("ab,ac->bc", arrays[0], arrays[1])
+        left, right = _truncated_svd(theta, chi_out)
+        return [left.T, right]
+
+    # (b, u0, d0) x (b, u1, d1) -> (u0, d0, u1, d1)
+    theta = contract("aij,akl->ijkl", arrays[0], arrays[1])
+    up_l, down_l, up_r, down_r = theta.shape
+    left, right = _truncated_svd(theta.reshape(up_l * down_l, up_r * down_r), chi_out)
+    rank = left.shape[1]
+    return [
+        left.reshape(up_l, down_l, rank).transpose(2, 0, 1),
+        right.reshape(rank, up_r, down_r),
+    ]
+
+
+def exact_apply(
+    left_tensor: Sequence[NDArray],
+    right_tensor: Sequence[NDArray],
+    chi_out: int,
+    kind: TrainKind,
+) -> list[NDArray]:
+    """Contract and compress two two-site trains exactly.
+
+    The MPO on the left is contracted site-wise with the right train, fusing
+    the two bond indices, and the result is compressed with a single SVD.
+
+    Args:
+        left_tensor: The two site tensors of the left MPO.
+        right_tensor: The two site tensors of the right MPS or MPO.
+        chi_out: The maximum bond dimension to keep.
+        kind: Whether ``right_tensor`` is an ``"mps"`` or an ``"mpo"``.
+
+    Returns:
+        The compressed product, in right-canonical form.
+
+    Raises:
+        ValueError: If either train does not have exactly two sites.
+    """
+    _check_pair(left_tensor)
+    _check_pair(right_tensor)
+    if kind == "mps":
+        # Contract the MPO lower leg with the MPS physical leg, fusing both bonds.
+        product = [
+            contract("aij,bj->abi", left_tensor[i], right_tensor[i]).reshape(
+                -1, left_tensor[i].shape[1]
+            )
+            for i in range(2)
+        ]
+    else:
+        product = [
+            contract("aij,bjk->abik", left_tensor[i], right_tensor[i]).reshape(
+                -1, left_tensor[i].shape[1], right_tensor[i].shape[2]
+            )
+            for i in range(2)
+        ]
+    return exact_compress(product, chi_out, kind)
+
+
+def _truncated_svd(theta: NDArray, chi_out: int) -> tuple[NDArray, NDArray]:
+    """Split a matrix as ``(U @ diag(S), Vh)``, keeping at most ``chi_out`` values."""
+    U, S, Vh = np.linalg.svd(theta, full_matrices=False)
+    rank = min(chi_out, S.size)
+    return U[:, :rank] * S[:rank], Vh[:rank]
+
+
+def _check_pair(arrays: Sequence[NDArray]) -> None:
+    """Reject trains that the exact two-site path cannot handle.
+
+    Raises:
+        ValueError: If the train does not have exactly two sites.
+    """
+    if len(arrays) != 2:
+        msg = (
+            f"Expected a two-site tensor train, got {len(arrays)} site(s). "
+            "Single-site trains are degenerate; use three or more sites for SRC."
+        )
+        raise ValueError(msg)
diff --git a/src/src_method/apply.py b/src/src_method/apply.py
index 4c864c3..433525f 100644
--- a/src/src_method/apply.py
+++ b/src/src_method/apply.py
@@ -14,10 +14,10 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-import quimb.tensor as qtn
 import structlog
 from opt_einsum import contract
 
+from ._tensor_train import MIN_SRC_SITES, exact_apply, infer_kind
 from .utils import (
     default_rng,
     get_xp,
@@ -27,8 +27,11 @@
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
     from types import ModuleType
 
+    from numpy.typing import NDArray
+
 # Set up logger
 setup_logging()
 logger = structlog.get_logger(__name__)
@@ -39,7 +42,7 @@
 LOG_TIME = " - Elapsed time (s)"
 LOG_WARN_SMALL = (
     "The current SRC implementation targets tensor networks with 3 or more sites. "
-    "Defaulting to the corresponding quimb primitive with SVD."
+    "Defaulting to an exact SVD-based contraction-compression."
 )
 
 # -----------------------------------------------
@@ -48,26 +51,28 @@
 
 
 def apply(
-    left_tensor: qtn.MatrixProductOperator,
-    right_tensor: qtn.MatrixProductOperator | qtn.MatrixProductState,
+    left_tensor: Sequence[NDArray],
+    right_tensor: Sequence[NDArray],
     chi_out: int,
     *,
     cutoff: float = 0.0,
     dtype: type = np.float64,
     seed: int | None = None,
     device: str = "cpu",
-) -> qtn.MatrixProductState | qtn.MatrixProductOperator:
+) -> list[NDArray]:
     """Applies the Successive Randomized Compression (SRC) algorithm.
 
-    Dispatches to the appropriate implementation based on the types of the input tensors.
-    Supported combinations are:
+    Tensor trains are plain lists of per-site arrays following the default
+    `quimb` index ordering; see `src_method._tensor_train` for the layout.
+    The train type is inferred from the rank of the first site tensor, and
+    dispatch follows:
 
       1. MPO-MPS: `left_tensor` is an MPO and `right_tensor` is an MPS. Results in an MPS.
       2. MPO-MPO: both `left_tensor` and `right_tensor` are MPOs. Results in an MPO.
 
     Args:
-        left_tensor: The left tensor network (MPO).
-        right_tensor: The right tensor network (MPO or MPS).
+        left_tensor: The site arrays of the left tensor network (MPO).
+        right_tensor: The site arrays of the right tensor network (MPO or MPS).
         chi_out: The desired maximum bond dimension of the output tensor network.
         cutoff: Relative singular-value cutoff for adaptive bond truncation.
             When positive, bonds are trimmed to their effective rank by
@@ -81,7 +86,7 @@ def apply(
             the optional ``cupy`` dependency for GPU execution.
 
     Returns:
-        The resulting compressed tensor network (MPS or MPO).
+        The site arrays of the compressed tensor network (MPS or MPO).
 
     Raises:
         TypeError: If the combination of input tensor types is unsupported.
@@ -91,28 +96,26 @@ def apply(
     xp = get_xp(device)
     prng = default_rng(seed)
 
-    if left_tensor.nsites < 3:
-        logger.warning(LOG_WARN_SMALL)
-        return left_tensor.apply(
-            right_tensor, compress=True, max_bond=chi_out, method="svd"
+    left_kind = infer_kind(left_tensor)
+    right_kind = infer_kind(right_tensor)
+    if left_kind != "mpo" or right_kind is None:
+        msg = (
+            "Unsupported combination of tensor network types: "
+            f"{left_kind or 'unknown'} and {right_kind or 'unknown'}; "
+            "expected an MPO on the left and an MPS or MPO on the right."
         )
-    if isinstance(left_tensor, qtn.MatrixProductOperator) and isinstance(
-        right_tensor, qtn.MatrixProductState
-    ):
+        raise TypeError(msg)
+
+    if len(left_tensor) < MIN_SRC_SITES:
+        logger.warning(LOG_WARN_SMALL)
+        return exact_apply(left_tensor, right_tensor, chi_out, right_kind)
+    if right_kind == "mps":
         return _src_mpo_mps(
             left_tensor, right_tensor, chi_out, prng, xp, cutoff=cutoff, dtype=dtype
         )
-    if isinstance(left_tensor, qtn.MatrixProductOperator) and isinstance(
-        right_tensor, qtn.MatrixProductOperator
-    ):
-        return _src_mpo_mpo(
-            left_tensor, right_tensor, chi_out, prng, xp, cutoff=cutoff, dtype=dtype
-        )
-    msg = (
-        "Unsupported combination of tensor network types: "
-        f"{type(left_tensor)} and {type(right_tensor)}"
+    return _src_mpo_mpo(
+        left_tensor, right_tensor, chi_out, prng, xp, cutoff=cutoff, dtype=dtype
     )
-    raise TypeError(msg)
 
 
 # -----------------------------------------------
@@ -121,20 +124,20 @@ def apply(
 
 
 def _src_mpo_mps(
-    mpo: qtn.MatrixProductOperator,
-    mps: qtn.MatrixProductState,
+    mpo: Sequence[NDArray],
+    mps: Sequence[NDArray],
     chi_out: int,
     prng: np.random.Generator,
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
     dtype: type = np.float64,
-) -> qtn.MatrixProductState:
+) -> list[NDArray]:
     """Computes the compressed product |η> ≈ H|ψ> using the SRC method.
 
     Args:
-        mpo: The MPO, as quimb object.
-        mps: The MPS, as a quimb object.
+        mpo: The site arrays of the MPO.
+        mps: The site arrays of the MPS.
         chi_out: The desired maximum bond dimension of the output MPS |η>.
         prng: A numpy / cupy random number generator.
         xp: Array module (``numpy`` or ``cupy``).
@@ -142,18 +145,18 @@ def _src_mpo_mps(
         dtype: The data type for the computation.
 
     Returns:
-        The resulting compressed MPS |η> in right-canonical form.
+        The site arrays of the compressed MPS |η> in right-canonical form.
     """
     # Problem dimensions
-    n_sites = mpo.nsites
-    phys_dim = mps[0].ind_size(mps.site_ind(0))
+    n_sites = len(mpo)
+    _, phys_dim = mps[0].shape
     logger.info(
         "Starting SRC MPO-MPS", n_sites=n_sites, phys_dim=phys_dim, device=xp.__name__
     )
 
     # Views of the tensors (transferred to device once, up front)
-    mpo_arrs = [xp.asarray(mpo[i].data) for i in range(n_sites)]
-    mps_arrs = [xp.asarray(mps[i].data) for i in range(n_sites)]
+    mpo_arrs = [xp.asarray(arr) for arr in mpo]
+    mps_arrs = [xp.asarray(arr) for arr in mps]
 
     # ----------------------------------------------
     # --- Left-to-Right Sweep: Compute C tensors ---
@@ -230,24 +233,24 @@ def _src_mpo_mps(
     logger.debug(LOG_TIME, t_rtl=tms * 1e-9)
     logger.info("SRC MPO-MPS complete.")
 
-    return qtn.MatrixProductState([to_numpy(t) for t in eta])
+    return [to_numpy(t) for t in eta]
 
 
 def _src_mpo_mpo(
-    mpo_left: qtn.MatrixProductOperator,
-    mpo_right: qtn.MatrixProductOperator,
+    mpo_left: Sequence[NDArray],
+    mpo_right: Sequence[NDArray],
     chi_out: int,
     prng: np.random.Generator,
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
     dtype: type = np.float64,
-) -> qtn.MatrixProductOperator:
+) -> list[NDArray]:
     """Computes the compressed product H_new ≈ H1 @ H2 using the SRC method.
 
     Args:
-        mpo_left: The first MPO.
-        mpo_right: The second MPO.
+        mpo_left: The site arrays of the first MPO.
+        mpo_right: The site arrays of the second MPO.
         chi_out: The desired maximum bond dimension of the output MPO.
         prng: A numpy / cupy random number generator.
         xp: Array module (``numpy`` or ``cupy``).
@@ -255,11 +258,11 @@ def _src_mpo_mpo(
         dtype: The data type for the computation.
 
     Returns:
-        The resulting compressed MPO in right-canonical form.
+        The site arrays of the compressed MPO in right-canonical form.
     """
     # Problem dimensions
-    n_sites = mpo_left.nsites
-    phys_up, phys_down = mpo_left.arrays[0].shape[1], mpo_right.arrays[0].shape[2]
+    n_sites = len(mpo_left)
+    phys_up, phys_down = mpo_left[0].shape[1], mpo_right[0].shape[2]
     logger.info(
         "Starting SRC MPO-MPO",
         n_sites=n_sites,
@@ -269,8 +272,8 @@ def _src_mpo_mpo(
     )
 
     # Views of the tensors (transferred to device once, up front)
-    mpo_left_arrs = [xp.asarray(mpo_left[i].data) for i in range(n_sites)]
-    mpo_right_arrs = [xp.asarray(mpo_right[i].data) for i in range(n_sites)]
+    mpo_left_arrs = [xp.asarray(arr) for arr in mpo_left]
+    mpo_right_arrs = [xp.asarray(arr) for arr in mpo_right]
 
     # ----------------------------------------------
     # --- Left-to-Right Sweep: Compute C tensors ---
@@ -351,4 +354,4 @@ def _src_mpo_mpo(
     logger.debug(LOG_TIME, t_rtl=tms * 1e-9)
     logger.info("SRC MPO-MPO complete.")
 
-    return qtn.MatrixProductOperator([to_numpy(t) for t in eta])
+    return [to_numpy(t) for t in eta]
diff --git a/src/src_method/compress.py b/src/src_method/compress.py
index 9f24919..b41abab 100644
--- a/src/src_method/compress.py
+++ b/src/src_method/compress.py
@@ -14,10 +14,10 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
-import quimb.tensor as qtn
 import structlog
 from opt_einsum import contract
 
+from ._tensor_train import MIN_SRC_SITES, exact_compress, infer_kind
 from .utils import (
     default_rng,
     get_xp,
@@ -27,8 +27,11 @@
 )
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
     from types import ModuleType
 
+    from numpy.typing import NDArray
+
 # Set up logger
 setup_logging()
 logger = structlog.get_logger(__name__)
@@ -39,7 +42,7 @@
 LOG_TIME = " - Elapsed time (s)"
 LOG_WARN_SMALL = (
     "The current SRC implementation targets tensor networks with 3 or more sites. "
-    "Defaulting to the corresponding quimb primitive with SVD."
+    "Defaulting to an exact SVD-based compression."
 )
 
 # -----------------------------------------------
@@ -48,24 +51,25 @@
 
 
 def compress(
-    tensor: qtn.MatrixProductState | qtn.MatrixProductOperator,
+    tensor: Sequence[NDArray],
     chi_out: int,
     *,
     cutoff: float = 0.0,
     dtype: type = np.float64,
     seed: int | None = None,
     device: str = "cpu",
-) -> qtn.MatrixProductState | qtn.MatrixProductOperator:
+) -> list[NDArray]:
     """Applies the Successive Randomized Compression (SRC) algorithm.
 
-    Dispatches to the appropriate implementation based on the types of the input tensors.
-    Supported combinations are:
+    Tensor trains are plain lists of per-site arrays following the default
+    `quimb` index ordering; see `src_method._tensor_train` for the layout.
+    The train type is inferred from the rank of the first site tensor:
 
       1. MPS: `tensor` is an MPS. Results in an MPS.
       2. MPO: `tensor` is an MPO. Results in an MPO.
 
     Args:
-        tensor: The tensor network to compress (MPS or MPO).
+        tensor: The site arrays of the tensor network to compress (MPS or MPO).
         chi_out: The desired maximum bond dimension of the output tensor network.
         cutoff: Relative singular-value cutoff for adaptive bond truncation.
             When positive, bonds are trimmed to their effective rank by
@@ -79,7 +83,7 @@ def compress(
             the optional ``cupy`` dependency for GPU execution.
 
     Returns:
-        The resulting compressed tensor network (MPS or MPO).
+        The site arrays of the compressed tensor network (MPS or MPO).
 
     Raises:
         TypeError: If the input tensor type is unsupported.
@@ -90,17 +94,20 @@ def compress(
     xp = get_xp(device)
     prng = default_rng(seed)
 
-    if tensor.nsites < 3:
+    kind = infer_kind(tensor)
+    if kind is None:
+        msg = (
+            "Unsupported tensor network layout: expected an MPS or MPO given as a "
+            "list of per-site arrays."
+        )
+        raise TypeError(msg)
+
+    if len(tensor) < MIN_SRC_SITES:
         logger.warning(LOG_WARN_SMALL)
-        compressed_tensor = tensor.copy(deep=True)
-        compressed_tensor.compress(max_bond=chi_out, method="svd")
-        return compressed_tensor
-    if isinstance(tensor, qtn.MatrixProductState):
+        return exact_compress(tensor, chi_out, kind)
+    if kind == "mps":
         return _src_mps(tensor, chi_out, prng, xp, cutoff=cutoff, dtype=dtype)
-    if isinstance(tensor, qtn.MatrixProductOperator):
-        return _src_mpo(tensor, chi_out, prng, xp, cutoff=cutoff, dtype=dtype)
-    msg = f"Unsupported tensor network type: {type(tensor)}"
-    raise TypeError(msg)
+    return _src_mpo(tensor, chi_out, prng, xp, cutoff=cutoff, dtype=dtype)
 
 
 # -----------------------------------------------
@@ -109,18 +116,18 @@ def compress(
 
 
 def _src_mpo(
-    mpo: qtn.MatrixProductOperator,
+    mpo: Sequence[NDArray],
     chi_out: int,
     prng: np.random.Generator,
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
     dtype: type = np.float64,
-) -> qtn.MatrixProductOperator:
+) -> list[NDArray]:
     """Compress an MPO using the SRC method.
 
     Args:
-        mpo: The MPO to compress as a quimb object.
+        mpo: The site arrays of the MPO to compress.
         chi_out: The desired maximum bond dimension of the output MPO.
         prng: A numpy / cupy random number generator instance.
         xp: Array module (``numpy`` or ``cupy``).
@@ -128,11 +135,11 @@ def _src_mpo(
         dtype: The data type for the computation.
 
     Returns:
-        The resulting compressed MPO as quimb object.
+        The site arrays of the compressed MPO.
     """
     # Problem dimensions
-    n_sites = mpo.nsites
-    _, phys_up, phys_down = mpo.arrays[0].shape
+    n_sites = len(mpo)
+    _, phys_up, phys_down = mpo[0].shape
     logger.info(
         "Starting SRC MPO",
         n_sites=n_sites,
@@ -142,7 +149,7 @@ def _src_mpo(
     )
 
     # Views of the tensors (transferred to device once, up front)
-    mpo_arrs = [xp.asarray(mpo[i].data) for i in range(n_sites)]
+    mpo_arrs = [xp.asarray(arr) for arr in mpo]
 
     # ----------------------------------------------
     # --- Left-to-Right Sweep: Compute C tensors ---
@@ -209,22 +216,22 @@ def _src_mpo(
     logger.debug(LOG_TIME, t_rtl=tms * 1e-9)
     logger.info("SRC MPO complete.")
 
-    return qtn.MatrixProductOperator([to_numpy(t) for t in eta])
+    return [to_numpy(t) for t in eta]
 
 
 def _src_mps(
-    mps: qtn.MatrixProductState,
+    mps: Sequence[NDArray],
     chi_out: int,
     prng: np.random.Generator,
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
     dtype: type = np.float64,
-) -> qtn.MatrixProductState:
+) -> list[NDArray]:
     """Compress an MPS using the SRC method.
 
     Args:
-        mps: The MPS to compress as a quimb object.
+        mps: The site arrays of the MPS to compress.
         chi_out: The desired maximum bond dimension of the output MPS |η>.
         prng: A numpy / cupy random number generator instance.
         xp: Array module (``numpy`` or ``cupy``).
@@ -232,17 +239,17 @@ def _src_mps(
         dtype: The data type for the computation.
 
     Returns:
-        The resulting compressed MPS as a quimb object.
+        The site arrays of the compressed MPS.
     """
     # Problem dimensions
-    n_sites = mps.nsites
-    _, phys_dim = mps.arrays[0].shape
+    n_sites = len(mps)
+    _, phys_dim = mps[0].shape
     logger.info(
         "Starting SRC MPS", n_sites=n_sites, phys_dim=phys_dim, device=xp.__name__
     )
 
     # View of the tensors (transferred to device once, up front)
-    mps_arrs = [xp.asarray(mps[i].data) for i in range(n_sites)]
+    mps_arrs = [xp.asarray(arr) for arr in mps]
 
     # ----------------------------------------------
     # --- Left-to-Right Sweep: Compute C tensors ---
@@ -303,4 +310,4 @@ def _src_mps(
     logger.debug(LOG_TIME, t_rtl=tms * 1e-9)
     logger.info("SRC MPS complete.")
 
-    return qtn.MatrixProductState([to_numpy(t) for t in eta])
+    return [to_numpy(t) for t in eta]
diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py
index be80600..2519538 100644
--- a/tests/test_gpu_backend.py
+++ b/tests/test_gpu_backend.py
@@ -14,6 +14,16 @@
 cupy = pytest.importorskip("cupy")
 
 
+def as_mps(arrays: list[np.ndarray]) -> qtn.MatrixProductState:
+    """Wrap a list of site arrays returned by src_method into a quimb MPS."""
+    return qtn.MatrixProductState(arrays)
+
+
+def as_mpo(arrays: list[np.ndarray]) -> qtn.MatrixProductOperator:
+    """Wrap a list of site arrays returned by src_method into a quimb MPO."""
+    return qtn.MatrixProductOperator(arrays)
+
+
 @pytest.fixture(autouse=True)
 def _require_device() -> None:
     """Skip the whole module when no GPU runtime is reachable."""
@@ -32,7 +42,16 @@ def test_apply_mpo_mps_gpu_matches_cpu(device: str) -> None:
         n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128
     )
 
-    out = apply(H, psi, chi_out=chi_out, dtype=np.complex128, seed=0, device=device)
+    out = as_mps(
+        apply(
+            H.arrays,
+            psi.arrays,
+            chi_out=chi_out,
+            dtype=np.complex128,
+            seed=0,
+            device=device,
+        )
+    )
     ref = H.apply(psi, compress=False)
 
     np.testing.assert_allclose(ref.distance(out), 0.0, atol=1e-6)
@@ -49,7 +68,16 @@ def test_apply_mpo_mpo_gpu_matches_cpu(device: str) -> None:
         n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128
     )
 
-    out = apply(H1, H2, chi_out=chi_out, dtype=np.complex128, seed=0, device=device)
+    out = as_mpo(
+        apply(
+            H1.arrays,
+            H2.arrays,
+            chi_out=chi_out,
+            dtype=np.complex128,
+            seed=0,
+            device=device,
+        )
+    )
     ref = H1.apply(H2, compress=False)
 
     np.testing.assert_allclose(ref.distance(out), 0.0, atol=1e-6)
@@ -63,7 +91,9 @@ def test_compress_mpo_gpu_matches_cpu(device: str) -> None:
     B = qtn.MPO_rand(n_sites, bond_dim=5, phys_dim=phys_dim, dtype=np.complex128) / 1e8
     C = A + B
 
-    out = compress(C, chi_out=chi_out, dtype=np.complex128, seed=0, device=device)
+    out = as_mpo(
+        compress(C.arrays, chi_out=chi_out, dtype=np.complex128, seed=0, device=device)
+    )
 
     np.testing.assert_allclose(C.distance(out), 0.0, atol=1e-6)
 
@@ -74,8 +104,26 @@ def test_apply_mpo_mpo_cpu_gpu_equivalent() -> None:
     H1 = qtn.MPO_rand(n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128)
     H2 = qtn.MPO_rand(n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128)
 
-    cpu_out = apply(H1, H2, chi_out=chi_out, dtype=np.complex128, seed=42, device="cpu")
-    gpu_out = apply(H1, H2, chi_out=chi_out, dtype=np.complex128, seed=42, device="gpu")
+    cpu_out = as_mpo(
+        apply(
+            H1.arrays,
+            H2.arrays,
+            chi_out=chi_out,
+            dtype=np.complex128,
+            seed=42,
+            device="cpu",
+        )
+    )
+    gpu_out = as_mpo(
+        apply(
+            H1.arrays,
+            H2.arrays,
+            chi_out=chi_out,
+            dtype=np.complex128,
+            seed=42,
+            device="gpu",
+        )
+    )
 
     # Gauge freedom means individual tensors can differ; compare the contracted MPOs.
     # atol=1e-6 accounts for floating-point accumulation across different execution orders.
@@ -87,8 +135,12 @@ def test_compress_mpo_cpu_gpu_equivalent() -> None:
     n_sites, phys_dim, chi_out = 8, 4, 32
     A = qtn.MPO_rand(n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128)
 
-    cpu_out = compress(A, chi_out=chi_out, dtype=np.complex128, seed=42, device="cpu")
-    gpu_out = compress(A, chi_out=chi_out, dtype=np.complex128, seed=42, device="gpu")
+    cpu_out = as_mpo(
+        compress(A.arrays, chi_out=chi_out, dtype=np.complex128, seed=42, device="cpu")
+    )
+    gpu_out = as_mpo(
+        compress(A.arrays, chi_out=chi_out, dtype=np.complex128, seed=42, device="gpu")
+    )
 
     np.testing.assert_allclose(cpu_out.distance(gpu_out), 0.0, atol=1e-6)
 
@@ -98,4 +150,4 @@ def test_invalid_device_raises() -> None:
     H = qtn.MPO_rand(5, bond_dim=4, phys_dim=2, dtype=np.complex128)
     psi = qtn.MPS_rand_state(5, bond_dim=4, phys_dim=2, dtype=np.complex128)
     with pytest.raises(ValueError, match="Unknown device"):
-        apply(H, psi, chi_out=4, device="tpu")
+        apply(H.arrays, psi.arrays, chi_out=4, device="tpu")
diff --git a/tests/test_package.py b/tests/test_package.py
index 983aa7f..aa28cfc 100644
--- a/tests/test_package.py
+++ b/tests/test_package.py
@@ -18,6 +18,20 @@
 # --- Utils ---
 # -------------
 
+# ``src_method`` operates on plain lists of site arrays.  quimb is used here only
+# to build reference networks and to measure distances, so every call unwraps its
+# inputs with ``.arrays`` and re-wraps the result with one of the helpers below.
+
+
+def as_mps(arrays: list[np.ndarray]) -> qtn.MatrixProductState:
+    """Wrap a list of site arrays returned by src_method into a quimb MPS."""
+    return qtn.MatrixProductState(arrays)
+
+
+def as_mpo(arrays: list[np.ndarray]) -> qtn.MatrixProductOperator:
+    """Wrap a list of site arrays returned by src_method into a quimb MPO."""
+    return qtn.MatrixProductOperator(arrays)
+
 
 def random_mpo(
     bonds: list[int],
@@ -64,7 +78,7 @@ def n_sites():
 
 
 @pytest.fixture
-def n_sites_quimb():
+def n_sites_small():
     return 2
 
 
@@ -108,7 +122,9 @@ def test_src_mpo_mps(n_sites, phys_dim, chi_out, array_type):
     )
 
     # SRC MPS should be identical to the original
-    psi_compress = apply(H, psi, chi_out=chi_out, dtype=array_type)
+    psi_compress = as_mps(
+        apply(H.arrays, psi.arrays, chi_out=chi_out, dtype=array_type)
+    )
 
     np.testing.assert_allclose(psi.distance(psi_compress), 0.0, atol=1e-6)
 
@@ -121,10 +137,10 @@ def test_src_mpo_mps_trims_terminal_bond() -> None:
         n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128, seed=3
     )
 
-    psi_src = apply(H, psi, chi_out=chi_out, dtype=np.complex128, seed=0)
+    psi_src = apply(H.arrays, psi.arrays, chi_out=chi_out, dtype=np.complex128, seed=0)
 
-    assert psi_src.arrays[-1].shape[0] < chi_out
-    assert psi_src.arrays[-1].shape[0] <= phys_dim
+    assert psi_src[-1].shape[0] < chi_out
+    assert psi_src[-1].shape[0] <= phys_dim
 
 
 # --------------------------------------------
@@ -140,7 +156,7 @@ def test_src_mpo_mpo_identity(n_sites, phys_dim, chi_out, array_type):
     H2 = qtn.MPO_identity(n_sites, phys_dim=phys_dim, dtype=array_type)
 
     # The compressed product should be identical to the original MPO
-    H_compress = apply(H1, H2, chi_out=chi_out, dtype=array_type)
+    H_compress = as_mpo(apply(H1.arrays, H2.arrays, chi_out=chi_out, dtype=array_type))
 
     np.testing.assert_allclose(H1.distance(H_compress), 0.0, atol=1e-6)
 
@@ -160,7 +176,7 @@ def test_src_mpo_mpo(n_sites, phys_dim, chi_out, array_type):
     H_ref = H1.apply(H2, compress=False)
 
     # The compressed product should be identical to the original MPO
-    H_src = apply(H1, H2, chi_out=chi_out, dtype=array_type)
+    H_src = as_mpo(apply(H1.arrays, H2.arrays, chi_out=chi_out, dtype=array_type))
 
     np.testing.assert_allclose(H_ref.distance(H_src), 0.0, atol=1e-6)
 
@@ -183,7 +199,7 @@ def test_src_mpo_mpo_long_phys(n_sites, phys_dim, chi_out, array_type):
     H_ref = H1.apply(H2, compress=False)
 
     # The compressed product should be identical to the original MPO
-    H_src = apply(H1, H2, chi_out=chi_out, dtype=array_type)
+    H_src = as_mpo(apply(H1.arrays, H2.arrays, chi_out=chi_out, dtype=array_type))
 
     np.testing.assert_allclose(H_ref.distance(H_src), 0.0, atol=1e-6)
 
@@ -194,7 +210,14 @@ def test_src_mpo_mpo_jagged(mpo_jagged_left, mpo_jagged_right, array_type):
     H_ref = mpo_jagged_left.apply(mpo_jagged_right, compress=False)
 
     # The compressed product should be identical to the original MPO
-    H_src = apply(mpo_jagged_left, mpo_jagged_right, chi_out=100, dtype=array_type)
+    H_src = as_mpo(
+        apply(
+            mpo_jagged_left.arrays,
+            mpo_jagged_right.arrays,
+            chi_out=100,
+            dtype=array_type,
+        )
+    )
 
     np.testing.assert_allclose(H_ref.distance(H_src), 0.0, atol=1e-6)
 
@@ -211,10 +234,10 @@ def test_src_mpo_mpo_trims_terminal_bond() -> None:
         n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128, seed=2
     )
 
-    H_src = apply(H1, H2, chi_out=chi_out, dtype=np.complex128, seed=0)
+    H_src = apply(H1.arrays, H2.arrays, chi_out=chi_out, dtype=np.complex128, seed=0)
 
-    assert H_src.arrays[-1].shape[0] < chi_out
-    assert H_src.arrays[-1].shape[0] <= phys_dim**2
+    assert H_src[-1].shape[0] < chi_out
+    assert H_src[-1].shape[0] <= phys_dim**2
 
 
 # --------------------------------
@@ -231,7 +254,7 @@ def test_src_mpo_compression(n_sites, phys_dim, chi_out, array_type):
     C = A + B
 
     # SRC MPO. It should be trivially compressed.
-    D = compress(C, chi_out=chi_out, dtype=array_type)
+    D = as_mpo(compress(C.arrays, chi_out=chi_out, dtype=array_type))
 
     np.testing.assert_allclose(C.distance(D), 0.0, atol=1e-6)
 
@@ -250,10 +273,10 @@ def test_src_mpo_compression_trims_terminal_bond() -> None:
     )
     C = A + B
 
-    D = compress(C, chi_out=chi_out, dtype=np.complex128, seed=0)
+    D = compress(C.arrays, chi_out=chi_out, dtype=np.complex128, seed=0)
 
-    assert D.arrays[-1].shape[0] < chi_out
-    assert D.arrays[-1].shape[0] <= phys_dim**2
+    assert D[-1].shape[0] < chi_out
+    assert D[-1].shape[0] <= phys_dim**2
 
 
 # --------------------------------
@@ -275,7 +298,7 @@ def test_src_mps(n_sites, phys_dim, chi_out, array_type):
     C = A + B
 
     # SRC MPS. It should be trivially compressed.
-    D = compress(C, chi_out=chi_out, dtype=array_type)
+    D = as_mps(compress(C.arrays, chi_out=chi_out, dtype=array_type))
 
     np.testing.assert_allclose(C.distance(D), 0.0, atol=1e-6)
 
@@ -294,10 +317,10 @@ def test_src_mps_compression_trims_terminal_bond() -> None:
     )
     C = A + B
 
-    D = compress(C, chi_out=chi_out, dtype=np.complex128, seed=0)
+    D = compress(C.arrays, chi_out=chi_out, dtype=np.complex128, seed=0)
 
-    assert D.arrays[-1].shape[0] < chi_out
-    assert D.arrays[-1].shape[0] <= phys_dim
+    assert D[-1].shape[0] < chi_out
+    assert D[-1].shape[0] <= phys_dim
 
 
 def test_src_mps_small_chi(n_sites, phys_dim, chi_out, array_type):
@@ -318,52 +341,88 @@ def test_src_mps_small_chi(n_sites, phys_dim, chi_out, array_type):
     C = A + B
 
     # SRC MPS. It should be trivially compressed.
-    D = compress(C, chi_out=chi_small, dtype=array_type)
+    D = as_mps(compress(C.arrays, chi_out=chi_small, dtype=array_type))
 
     np.testing.assert_allclose(C.distance(D), 0.0, atol=1e-6)
 
 
-# -----------------------------------------------
-# --- Test default to quimb for small systems ---
-# -----------------------------------------------
+# ---------------------------------------------------
+# --- Test exact fallback for small (<3 site) TNs ---
+# ---------------------------------------------------
 
 
-def test_apply_quimb_dispatch(n_sites_quimb, phys_dim, chi_out, array_type):
-    """Tests that apply dispatches to quimb for small systems."""
+def test_apply_small_system_dispatch(n_sites_small, phys_dim, chi_out, array_type):
+    """Tests that apply falls back to the exact path for small systems."""
 
     # Generate a random MPS and connect it to the identity MPO
-    H = qtn.MPO_identity(n_sites_quimb, phys_dim=phys_dim, dtype=array_type)
+    H = qtn.MPO_identity(n_sites_small, phys_dim=phys_dim, dtype=array_type)
     psi = qtn.MPS_rand_state(
-        n_sites_quimb, bond_dim=chi_out, phys_dim=phys_dim, dtype=array_type
+        n_sites_small, bond_dim=chi_out, phys_dim=phys_dim, dtype=array_type
     )
 
     # SRC MPS should be identical to the original
-    psi_compress = apply(H, psi, chi_out=chi_out, dtype=array_type)
+    psi_compress = as_mps(
+        apply(H.arrays, psi.arrays, chi_out=chi_out, dtype=array_type)
+    )
 
     np.testing.assert_allclose(psi.distance(psi_compress), 0.0, atol=1e-6)
 
 
-def test_compress_quimb_dispatch(n_sites_quimb, phys_dim, chi_out, array_type):
-    """Tests that compress dispatches to quimb for small systems."""
+def test_apply_small_system_mpo_mpo(n_sites_small, phys_dim, chi_out, array_type):
+    """The exact fallback must also handle two-site MPO-MPO products."""
+    H1 = qtn.MPO_rand(
+        n_sites_small, bond_dim=chi_out, phys_dim=phys_dim, dtype=array_type
+    )
+    H2 = qtn.MPO_identity(n_sites_small, phys_dim=phys_dim, dtype=array_type)
+
+    H_src = as_mpo(apply(H1.arrays, H2.arrays, chi_out=chi_out, dtype=array_type))
+
+    np.testing.assert_allclose(H1.distance(H_src), 0.0, atol=1e-6)
+
+
+def test_compress_small_system_dispatch(n_sites_small, phys_dim, chi_out, array_type):
+    """Tests that compress falls back to the exact path for small systems."""
 
     # Add an almost-zero MPS B to a dense MPS A, inflating the bond dimension
     A = qtn.MPS_rand_state(
-        n_sites_quimb, bond_dim=chi_out, phys_dim=phys_dim, dtype=array_type
+        n_sites_small, bond_dim=chi_out, phys_dim=phys_dim, dtype=array_type
     )
     B = (
         qtn.MPS_rand_state(
-            n_sites_quimb, bond_dim=5, phys_dim=phys_dim, dtype=array_type
+            n_sites_small, bond_dim=5, phys_dim=phys_dim, dtype=array_type
         )
         / 1e8
     )
     C = A + B
 
     # SRC MPS. It should be trivially compressed.
-    D = compress(C, chi_out=chi_out, dtype=array_type)
+    D = as_mps(compress(C.arrays, chi_out=chi_out, dtype=array_type))
+
+    np.testing.assert_allclose(C.distance(D), 0.0, atol=1e-6)
+
+
+def test_compress_small_system_mpo(n_sites_small, phys_dim, chi_out, array_type):
+    """The exact fallback must also handle two-site MPOs."""
+    A = qtn.MPO_rand(
+        n_sites_small, bond_dim=chi_out, phys_dim=phys_dim, dtype=array_type
+    )
+    B = (
+        qtn.MPO_rand(n_sites_small, bond_dim=2, phys_dim=phys_dim, dtype=array_type)
+        / 1e8
+    )
+    C = A + B
+
+    D = as_mpo(compress(C.arrays, chi_out=chi_out, dtype=array_type))
 
     np.testing.assert_allclose(C.distance(D), 0.0, atol=1e-6)
 
 
+def test_single_site_train_raises(phys_dim, chi_out, array_type):
+    """A single-site train is degenerate and must be rejected."""
+    with pytest.raises(ValueError, match="two-site tensor train"):
+        compress([np.ones((1, phys_dim), dtype=array_type)], chi_out=chi_out)
+
+
 # -------------------------------------------------
 # --- Test adaptive bond truncation via cutoff  ---
 # -------------------------------------------------
@@ -378,8 +437,14 @@ def test_cutoff_zero_matches_default_mpo_mpo(n_sites, phys_dim, chi_out, array_t
         n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=array_type
     )
 
-    H_default = apply(H1, H2, chi_out=chi_out, dtype=array_type, seed=42)
-    H_cutoff0 = apply(H1, H2, chi_out=chi_out, cutoff=0.0, dtype=array_type, seed=42)
+    H_default = as_mpo(
+        apply(H1.arrays, H2.arrays, chi_out=chi_out, dtype=array_type, seed=42)
+    )
+    H_cutoff0 = as_mpo(
+        apply(
+            H1.arrays, H2.arrays, chi_out=chi_out, cutoff=0.0, dtype=array_type, seed=42
+        )
+    )
 
     np.testing.assert_allclose(H_default.distance(H_cutoff0), 0.0, atol=1e-6)
 
@@ -390,8 +455,19 @@ def test_cutoff_trims_bonds_mpo_mpo(n_sites, phys_dim, chi_out, array_type):
     H2 = qtn.MPO_identity(n_sites, phys_dim=phys_dim, dtype=array_type)
 
     # Identity product: effective rank = chi_out, no truncation expected
-    H_no_cut = apply(H1, H2, chi_out=chi_out, dtype=array_type, seed=42)
-    H_cut = apply(H1, H2, chi_out=chi_out, cutoff=1e-10, dtype=array_type, seed=42)
+    H_no_cut = as_mpo(
+        apply(H1.arrays, H2.arrays, chi_out=chi_out, dtype=array_type, seed=42)
+    )
+    H_cut = as_mpo(
+        apply(
+            H1.arrays,
+            H2.arrays,
+            chi_out=chi_out,
+            cutoff=1e-10,
+            dtype=array_type,
+            seed=42,
+        )
+    )
 
     # Result should still be accurate
     np.testing.assert_allclose(H1.distance(H_cut), 0.0, atol=1e-5)
@@ -408,7 +484,9 @@ def test_cutoff_preserves_accuracy_mpo_mps(n_sites, phys_dim, chi_out, array_typ
         n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=array_type
     )
 
-    psi_cut = apply(H, psi, chi_out=chi_out, cutoff=1e-10, dtype=array_type)
+    psi_cut = as_mps(
+        apply(H.arrays, psi.arrays, chi_out=chi_out, cutoff=1e-10, dtype=array_type)
+    )
 
     np.testing.assert_allclose(psi.distance(psi_cut), 0.0, atol=1e-6)
 
@@ -419,7 +497,7 @@ def test_cutoff_preserves_accuracy_compress_mpo(n_sites, phys_dim, chi_out, arra
     B = qtn.MPO_rand(n_sites, bond_dim=5, phys_dim=phys_dim, dtype=array_type) / 1e8
     C = A + B
 
-    D = compress(C, chi_out=chi_out, cutoff=1e-10, dtype=array_type)
+    D = as_mpo(compress(C.arrays, chi_out=chi_out, cutoff=1e-10, dtype=array_type))
 
     np.testing.assert_allclose(C.distance(D), 0.0, atol=1e-6)
 
@@ -435,7 +513,7 @@ def test_cutoff_preserves_accuracy_compress_mps(n_sites, phys_dim, chi_out, arra
     )
     C = A + B
 
-    D = compress(C, chi_out=chi_out, cutoff=1e-10, dtype=array_type)
+    D = as_mps(compress(C.arrays, chi_out=chi_out, cutoff=1e-10, dtype=array_type))
 
     np.testing.assert_allclose(C.distance(D), 0.0, atol=1e-6)
 
@@ -458,18 +536,17 @@ def test_apply_unsupported_types(n_sites, phys_dim, chi_out, array_type):
 
     # Unsupported combination: MPS-MPS
     with pytest.raises(TypeError):
-        apply(psi, phi, chi_out=chi_out, dtype=array_type)
+        apply(psi.arrays, phi.arrays, chi_out=chi_out, dtype=array_type)
 
 
 def test_compress_unsupported_type(n_sites, chi_out, array_type):
-    """Tests that compress raises TypeError for unsupported tensor types."""
+    """Tests that compress raises TypeError for unsupported tensor layouts."""
 
-    # Generate a random PEPS, which is unsupported
-    peps = qtn.PEPS.rand(Lx=n_sites, Ly=n_sites, bond_dim=chi_out)
+    # Rank-4 boundary tensors are neither an MPS nor an MPO (e.g. a PEPS row)
+    peps_like = [np.zeros((2, 2, 2, 2), dtype=array_type)] * n_sites
 
-    # Unsupported type: PEPS
     with pytest.raises(TypeError):
-        compress(peps, chi_out=chi_out, dtype=array_type)
+        compress(peps_like, chi_out=chi_out, dtype=array_type)
 
 
 # -----------------------------------------
@@ -495,7 +572,9 @@ def test_benchmark_src_mpo_mpo(benchmark):
     )
 
     # Benchmark the application
-    result_mpo = benchmark(apply, H1, H2, chi_out=chi_out, dtype=array_type)
+    result_mpo = benchmark(
+        apply, H1.arrays, H2.arrays, chi_out=chi_out, dtype=array_type
+    )
 
     # Still has to be correct
-    np.testing.assert_allclose(H1.distance(result_mpo), 0.0, atol=1e-6)
+    np.testing.assert_allclose(H1.distance(as_mpo(result_mpo)), 0.0, atol=1e-6)
diff --git a/uv.lock b/uv.lock
index e1ea52b..c5b36d4 100644
--- a/uv.lock
+++ b/uv.lock
@@ -156,7 +156,7 @@ name = "cffi"
 version = "2.0.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "pycparser" },
+    { name = "pycparser", marker = "implementation_name != 'PyPy'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" }
 wheels = [
@@ -831,17 +831,17 @@ resolution-markers = [
     "python_full_version < '3.12'",
 ]
 dependencies = [
-    { name = "colorama", marker = "sys_platform == 'win32'" },
-    { name = "decorator" },
-    { name = "ipython-pygments-lexers" },
-    { name = "jedi" },
-    { name = "matplotlib-inline" },
-    { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
-    { name = "prompt-toolkit" },
-    { name = "pygments" },
-    { name = "stack-data" },
-    { name = "traitlets" },
-    { name = "typing-extensions" },
+    { name = "colorama", marker = "python_full_version < '3.12' and sys_platform == 'win32'" },
+    { name = "decorator", marker = "python_full_version < '3.12'" },
+    { name = "ipython-pygments-lexers", marker = "python_full_version < '3.12'" },
+    { name = "jedi", marker = "python_full_version < '3.12'" },
+    { name = "matplotlib-inline", marker = "python_full_version < '3.12'" },
+    { name = "pexpect", marker = "python_full_version < '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "prompt-toolkit", marker = "python_full_version < '3.12'" },
+    { name = "pygments", marker = "python_full_version < '3.12'" },
+    { name = "stack-data", marker = "python_full_version < '3.12'" },
+    { name = "traitlets", marker = "python_full_version < '3.12'" },
+    { name = "typing-extensions", marker = "python_full_version < '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/c5/25/daae0e764047b0a2480c7bbb25d48f4f509b5818636562eeac145d06dfee/ipython-9.10.1.tar.gz", hash = "sha256:e170e9b2a44312484415bdb750492699bf329233b03f2557a9692cce6466ada4", size = 4426663, upload-time = "2026-03-27T09:53:26.244Z" }
 wheels = [
@@ -856,16 +856,16 @@ resolution-markers = [
     "python_full_version >= '3.12'",
 ]
 dependencies = [
-    { name = "colorama", marker = "sys_platform == 'win32'" },
-    { name = "decorator" },
-    { name = "ipython-pygments-lexers" },
-    { name = "jedi" },
-    { name = "matplotlib-inline" },
-    { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
-    { name = "prompt-toolkit" },
-    { name = "pygments" },
-    { name = "stack-data" },
-    { name = "traitlets" },
+    { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" },
+    { name = "decorator", marker = "python_full_version >= '3.12'" },
+    { name = "ipython-pygments-lexers", marker = "python_full_version >= '3.12'" },
+    { name = "jedi", marker = "python_full_version >= '3.12'" },
+    { name = "matplotlib-inline", marker = "python_full_version >= '3.12'" },
+    { name = "pexpect", marker = "python_full_version >= '3.12' and sys_platform != 'emscripten' and sys_platform != 'win32'" },
+    { name = "prompt-toolkit", marker = "python_full_version >= '3.12'" },
+    { name = "pygments", marker = "python_full_version >= '3.12'" },
+    { name = "stack-data", marker = "python_full_version >= '3.12'" },
+    { name = "traitlets", marker = "python_full_version >= '3.12'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/3a/73/7114f80a8f9cabdb13c27732dce24af945b2923dcab80723602f7c8bc2d8/ipython-9.12.0.tar.gz", hash = "sha256:01daa83f504b693ba523b5a407246cabde4eb4513285a3c6acaff11a66735ee4", size = 4428879, upload-time = "2026-03-27T09:42:45.312Z" }
 wheels = [
@@ -2550,7 +2550,6 @@ dependencies = [
     { name = "numba" },
     { name = "numpy" },
     { name = "opt-einsum" },
-    { name = "quimb" },
     { name = "structlog" },
 ]
 
@@ -2575,6 +2574,7 @@ dev = [
     { name = "pytest" },
     { name = "pytest-benchmark" },
     { name = "pytest-cov" },
+    { name = "quimb" },
     { name = "ruff" },
     { name = "scalene" },
 ]
@@ -2604,6 +2604,7 @@ test = [
     { name = "pytest" },
     { name = "pytest-benchmark" },
     { name = "pytest-cov" },
+    { name = "quimb" },
 ]
 
 [package.metadata]
@@ -2623,7 +2624,6 @@ requires-dist = [
     { name = "nvidia-cusparse-cu12", marker = "extra == 'gpu-nvidia'" },
     { name = "nvidia-nvjitlink-cu12", marker = "extra == 'gpu-nvidia'" },
     { name = "opt-einsum", specifier = ">=3.4.0" },
-    { name = "quimb", specifier = ">=1.14.0,<2" },
     { name = "structlog", specifier = ">=25.5.0" },
 ]
 provides-extras = ["gpu-nvidia", "gpu-rocm"]
@@ -2634,6 +2634,7 @@ dev = [
     { name = "pytest", specifier = ">=9.1.0,<10.0" },
     { name = "pytest-benchmark", specifier = ">=5.2.3" },
     { name = "pytest-cov", specifier = ">=4.0" },
+    { name = "quimb", specifier = ">=1.14.0,<2" },
     { name = "ruff", specifier = "==0.16.5" },
     { name = "scalene" },
 ]
@@ -2663,6 +2664,7 @@ test = [
     { name = "pytest", specifier = ">=9.1.0,<10.0" },
     { name = "pytest-benchmark", specifier = ">=5.2.3" },
     { name = "pytest-cov", specifier = ">=4.0" },
+    { name = "quimb", specifier = ">=1.14.0,<2" },
 ]
 
 [[package]]

From 4c37390e8154a615db8837609ce137cadd097327 Mon Sep 17 00:00:00 2001
From: Panadestein <rpana92@gmail.com>
Date: Tue, 1 Sep 2026 19:47:48 +0000
Subject: [PATCH 2/4] =?UTF-8?q?build!:=20=F0=9F=94=A5=20drop=20unused=20ru?=
 =?UTF-8?q?ntime=20dependencies?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

With quimb gone, the remaining runtime dependency list was mostly dead
weight. `numba` and `llvmlite` are imported nowhere in `src/` — they were
transitively required by quimb, not by us. `cyclopts` is only used by the
benchmark scripts under `benches/`, which are not packaged.

`cmaes` is not dead: cotengra (via quimb) discovers it by name as a
hyper-optimization backend and warns when it is missing, which
`filterwarnings = ["error"]` promotes to a test failure. It moves to the
`test` group alongside quimb rather than being removed.

A runtime install is now numpy + opt_einsum + structlog.

- move `cyclopts` to the `dev` group (benchmarks)
- move `cmaes` to the `test` group (cotengra path optimizer)
- drop the now-dead `@jit(` / `@njit(` coverage excludes
---
 pyproject.toml | 20 ++++++++++++--------
 uv.lock        | 20 +++++++++-----------
 2 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 2a00f1e..b9042b3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,10 +32,6 @@ classifiers = [
 ]
 dynamic = ["version"]
 dependencies = [
-  "cmaes>=0.12.0",
-  "cyclopts>=3.20",
-  "llvmlite>=0.44",
-  "numba>=0.65.0",
   "numpy>=2.4.4,<3",
   "opt_einsum>=3.4.0",
   "structlog>=25.5.0",
@@ -68,13 +64,23 @@ Discussions = "https://github.com/Algorithmiq/src-method/discussions"
 [dependency-groups]
 # quimb is *not* a runtime dependency: it is only used to build reference
 # tensor networks and to measure distances in the tests and benchmarks.
+# cmaes is pulled in by quimb's path optimizer (cotengra), which warns when no
+# hyper-optimization backend is installed; `filterwarnings = ["error"]` would
+# otherwise turn that warning into a test failure.
 test = [
   "pytest-cov>=4.0",
   "pytest>=9.1.0,<10.0",
   "pytest-benchmark>=5.2.3",
   "quimb>=1.14.0,<2",
+  "cmaes>=0.12.0",
+]
+dev = [
+  "scalene",
+  "pre-commit>=3",
+  "ruff==0.16.5",
+  "cyclopts>=3.20",
+  { include-group = "test" },
 ]
-dev = ["scalene", "pre-commit>=3", "ruff==0.16.5", { include-group = "test" }]
 interactive = [
   "ipykernel>=7.2.0",
   "matplotlib>=3.11.0",
@@ -83,7 +89,7 @@ interactive = [
   "pylatexenc",
 ]
 docs = [
-  "black[jupyter]>=26.5.1",          # format extracted function signatures
+  "black[jupyter]>=26.5.1",
   "mkdocs-gen-files>=0.5.0",
   "mkdocs-jupyter>=0.24.2",
   "mkdocs-literate-nav>=0.6.0",
@@ -144,8 +150,6 @@ exclude_also = [
   "class .*\\bProtocol\\):",
   "@(abc\\.)?abstractmethod",
   "if (typing\\.)?TYPE_CHECKING:",
-  "@jit\\(",
-  "@njit\\(",
   '^\s*\.\.\.',
   "logger =",
   "logger\\.",
diff --git a/uv.lock b/uv.lock
index c5b36d4..4cfe14a 100644
--- a/uv.lock
+++ b/uv.lock
@@ -288,14 +288,14 @@ wheels = [
 
 [[package]]
 name = "cmaes"
-version = "0.13.0"
+version = "0.13.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "numpy" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/3d/9f/ae4edb7dec820e84fef7a90b753ae5c72c66a05ffa69a7894771024386a7/cmaes-0.13.0.tar.gz", hash = "sha256:69a252b0291d08100351e37c2918c7c6d929b02ab7dcd9dd14fc02c7c98cc1b9", size = 61265, upload-time = "2026-03-28T07:41:55.249Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/94/ce/dcf35c4db556fcd25547f62b3b3fbf9e69f44054ff85434310702e29ac52/cmaes-0.13.1.tar.gz", hash = "sha256:5f56a9e8dd769ffd23ac2ad651979bdb6ca8846de60d97d03fe5ad8456c56c79", size = 61781, upload-time = "2026-08-21T13:42:59.347Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/7e/98/be3f668f77838b2756ccc78a45e0c62f43d3134003f3f4bad814d37df1b3/cmaes-0.13.0-py3-none-any.whl", hash = "sha256:ccf61c73d5792cf44b50672b63f28c590082d1c1f5ab5155e2dd6e5305427cad", size = 73027, upload-time = "2026-03-28T07:41:53.956Z" },
+    { url = "https://files.pythonhosted.org/packages/ff/18/1004ca44ba64b8b5ee24968246e056ff2110916250cc1a1ac4239ea407a8/cmaes-0.13.1-py3-none-any.whl", hash = "sha256:659881b1d574adf0b378448821e3da24a4f7b7a1989640cd902000c567897db5", size = 73401, upload-time = "2026-08-21T13:42:57.941Z" },
 ]
 
 [[package]]
@@ -2544,10 +2544,6 @@ wheels = [
 name = "src-method"
 source = { editable = "." }
 dependencies = [
-    { name = "cmaes" },
-    { name = "cyclopts" },
-    { name = "llvmlite" },
-    { name = "numba" },
     { name = "numpy" },
     { name = "opt-einsum" },
     { name = "structlog" },
@@ -2570,6 +2566,8 @@ gpu-rocm = [
 
 [package.dev-dependencies]
 dev = [
+    { name = "cmaes" },
+    { name = "cyclopts" },
     { name = "pre-commit" },
     { name = "pytest" },
     { name = "pytest-benchmark" },
@@ -2601,6 +2599,7 @@ interactive = [
     { name = "tqdm" },
 ]
 test = [
+    { name = "cmaes" },
     { name = "pytest" },
     { name = "pytest-benchmark" },
     { name = "pytest-cov" },
@@ -2609,12 +2608,8 @@ test = [
 
 [package.metadata]
 requires-dist = [
-    { name = "cmaes", specifier = ">=0.12.0" },
     { name = "cupy", marker = "extra == 'gpu-rocm'", specifier = ">=14.1.1" },
     { name = "cupy-cuda12x", marker = "extra == 'gpu-nvidia'", specifier = ">=13" },
-    { name = "cyclopts", specifier = ">=3.20" },
-    { name = "llvmlite", specifier = ">=0.44" },
-    { name = "numba", specifier = ">=0.65.0" },
     { name = "numpy", specifier = ">=2.4.4,<3" },
     { name = "nvidia-cublas-cu12", marker = "extra == 'gpu-nvidia'" },
     { name = "nvidia-cuda-nvrtc-cu12", marker = "extra == 'gpu-nvidia'" },
@@ -2630,6 +2625,8 @@ provides-extras = ["gpu-nvidia", "gpu-rocm"]
 
 [package.metadata.requires-dev]
 dev = [
+    { name = "cmaes", specifier = ">=0.12.0" },
+    { name = "cyclopts", specifier = ">=3.20" },
     { name = "pre-commit", specifier = ">=3" },
     { name = "pytest", specifier = ">=9.1.0,<10.0" },
     { name = "pytest-benchmark", specifier = ">=5.2.3" },
@@ -2661,6 +2658,7 @@ interactive = [
     { name = "tqdm", specifier = ">=4.65.0" },
 ]
 test = [
+    { name = "cmaes", specifier = ">=0.12.0" },
     { name = "pytest", specifier = ">=9.1.0,<10.0" },
     { name = "pytest-benchmark", specifier = ">=5.2.3" },
     { name = "pytest-cov", specifier = ">=4.0" },

From 454c2ba882c17e4452c06f521035a75e24248fc8 Mon Sep 17 00:00:00 2001
From: Panadestein <rpana92@gmail.com>
Date: Tue, 1 Sep 2026 19:49:00 +0000
Subject: [PATCH 3/4] chore: remove verbose comments.

---
 pyproject.toml | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index b9042b3..540b06a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -62,11 +62,6 @@ Discussions = "https://github.com/Algorithmiq/src-method/discussions"
 
 
 [dependency-groups]
-# quimb is *not* a runtime dependency: it is only used to build reference
-# tensor networks and to measure distances in the tests and benchmarks.
-# cmaes is pulled in by quimb's path optimizer (cotengra), which warns when no
-# hyper-optimization backend is installed; `filterwarnings = ["error"]` would
-# otherwise turn that warning into a test failure.
 test = [
   "pytest-cov>=4.0",
   "pytest>=9.1.0,<10.0",

From dcd3b474b50a0dcf14e05a27b8e7a2a06d1e1a2e Mon Sep 17 00:00:00 2001
From: Panadestein <rpana92@gmail.com>
Date: Wed, 2 Sep 2026 08:51:55 +0000
Subject: [PATCH 4/4] =?UTF-8?q?fix:=20=F0=9F=90=9B=20validate=20site=20cou?=
 =?UTF-8?q?nts=20and=20accept=20device=20arrays=20in=20the=20exact=20path?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Addresses the Copilot review on #24.

The exact two-site path called `np.linalg.svd` on whatever the caller
passed, so it broke on device arrays while the SRC sweep handled them via
`xp.asarray`. It now brings inputs onto the host with `to_numpy`, matching
the sweep and always returning numpy arrays.

`apply` also accepted trains of different lengths: the sweep sizes itself
from the left train, so a 4-site MPO applied to a 6-site MPS silently
returned a 4-site result. quimb used to reject this; the list-based API
has to do it itself.

- reject mismatched site counts in `apply`
- move the two-site guard to the public boundary as `check_exact_supported`,
  so degenerate trains raise instead of first logging a fallback warning
- cover both guards, and assert the warning is not emitted when raising
---
 src/src_method/_tensor_train.py | 60 +++++++++++++++++++--------------
 src/src_method/apply.py         | 19 +++++++++--
 src/src_method/compress.py      | 11 ++++--
 tests/test_package.py           | 22 ++++++++++++
 4 files changed, 83 insertions(+), 29 deletions(-)

diff --git a/src/src_method/_tensor_train.py b/src/src_method/_tensor_train.py
index 37510da..99a101e 100644
--- a/src/src_method/_tensor_train.py
+++ b/src/src_method/_tensor_train.py
@@ -21,6 +21,8 @@
 import numpy as np
 from opt_einsum import contract
 
+from .utils import to_numpy
+
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
@@ -29,6 +31,9 @@
 # Minimum number of sites for which the randomized SRC sweep is defined.
 MIN_SRC_SITES = 3
 
+# The only sub-``MIN_SRC_SITES`` size the exact path can handle.
+_EXACT_SITES = 2
+
 # Rank of a boundary (first / last) site tensor, which identifies the train type.
 _MPS_BOUNDARY_NDIM = 2
 _MPO_BOUNDARY_NDIM = 3
@@ -38,6 +43,7 @@
 __all__ = [
     "MIN_SRC_SITES",
     "TrainKind",
+    "check_exact_supported",
     "exact_apply",
     "exact_compress",
     "infer_kind",
@@ -66,23 +72,43 @@ def infer_kind(arrays: Sequence[NDArray]) -> TrainKind | None:
     return None
 
 
+def check_exact_supported(n_sites: int) -> None:
+    """Reject sub-``MIN_SRC_SITES`` trains the exact path cannot handle.
+
+    Called at the public boundary before the fallback is announced, so that a
+    degenerate train raises instead of first logging a misleading warning.
+
+    Args:
+        n_sites: The number of sites in the train.
+
+    Raises:
+        ValueError: If the train does not have exactly two sites.
+    """
+    if n_sites != _EXACT_SITES:
+        msg = (
+            f"Expected a two-site tensor train, got {n_sites} site(s). "
+            "Single-site trains are degenerate; use three or more sites for SRC."
+        )
+        raise ValueError(msg)
+
+
 def exact_compress(
     arrays: Sequence[NDArray], chi_out: int, kind: TrainKind
 ) -> list[NDArray]:
     """Compress a two-site train exactly via a single truncated SVD.
 
+    Site counts are validated by the caller via `check_exact_supported`.
+
     Args:
         arrays: The two site tensors of the train.
         chi_out: The maximum bond dimension to keep.
         kind: Whether the train is an ``"mps"`` or an ``"mpo"``.
 
     Returns:
-        The compressed train, in right-canonical form.
-
-    Raises:
-        ValueError: If the train does not have exactly two sites.
+        The compressed train, in right-canonical form, as numpy arrays.
     """
-    _check_pair(arrays)
+    # The dense SVD is host-side, so accept device arrays like the sweep does.
+    arrays = [to_numpy(arr) for arr in arrays]
     if kind == "mps":
         # (b, p0) x (b, p1) -> (p0, p1)
         theta = contract("ab,ac->bc", arrays[0], arrays[1])
@@ -110,6 +136,7 @@ def exact_apply(
 
     The MPO on the left is contracted site-wise with the right train, fusing
     the two bond indices, and the result is compressed with a single SVD.
+    Site counts are validated by the caller via `check_exact_supported`.
 
     Args:
         left_tensor: The two site tensors of the left MPO.
@@ -118,13 +145,10 @@ def exact_apply(
         kind: Whether ``right_tensor`` is an ``"mps"`` or an ``"mpo"``.
 
     Returns:
-        The compressed product, in right-canonical form.
-
-    Raises:
-        ValueError: If either train does not have exactly two sites.
+        The compressed product, in right-canonical form, as numpy arrays.
     """
-    _check_pair(left_tensor)
-    _check_pair(right_tensor)
+    left_tensor = [to_numpy(arr) for arr in left_tensor]
+    right_tensor = [to_numpy(arr) for arr in right_tensor]
     if kind == "mps":
         # Contract the MPO lower leg with the MPS physical leg, fusing both bonds.
         product = [
@@ -148,17 +172,3 @@ def _truncated_svd(theta: NDArray, chi_out: int) -> tuple[NDArray, NDArray]:
     U, S, Vh = np.linalg.svd(theta, full_matrices=False)
     rank = min(chi_out, S.size)
     return U[:, :rank] * S[:rank], Vh[:rank]
-
-
-def _check_pair(arrays: Sequence[NDArray]) -> None:
-    """Reject trains that the exact two-site path cannot handle.
-
-    Raises:
-        ValueError: If the train does not have exactly two sites.
-    """
-    if len(arrays) != 2:
-        msg = (
-            f"Expected a two-site tensor train, got {len(arrays)} site(s). "
-            "Single-site trains are degenerate; use three or more sites for SRC."
-        )
-        raise ValueError(msg)
diff --git a/src/src_method/apply.py b/src/src_method/apply.py
index 433525f..4ea0e65 100644
--- a/src/src_method/apply.py
+++ b/src/src_method/apply.py
@@ -17,7 +17,12 @@
 import structlog
 from opt_einsum import contract
 
-from ._tensor_train import MIN_SRC_SITES, exact_apply, infer_kind
+from ._tensor_train import (
+    MIN_SRC_SITES,
+    check_exact_supported,
+    exact_apply,
+    infer_kind,
+)
 from .utils import (
     default_rng,
     get_xp,
@@ -90,7 +95,8 @@ def apply(
 
     Raises:
         TypeError: If the combination of input tensor types is unsupported.
-        ValueError: If ``device`` is not recognised.
+        ValueError: If the two trains differ in length, if a sub-three-site
+            train is not exactly two sites, or if ``device`` is not recognised.
         ImportError: If ``device="gpu"`` but cupy is not installed.
     """
     xp = get_xp(device)
@@ -106,7 +112,16 @@ def apply(
         )
         raise TypeError(msg)
 
+    # Without this the sweep would silently drop the extra sites of the longer train.
+    if len(left_tensor) != len(right_tensor):
+        msg = (
+            "Both tensor trains must have the same number of sites, got "
+            f"{len(left_tensor)} and {len(right_tensor)}."
+        )
+        raise ValueError(msg)
+
     if len(left_tensor) < MIN_SRC_SITES:
+        check_exact_supported(len(left_tensor))
         logger.warning(LOG_WARN_SMALL)
         return exact_apply(left_tensor, right_tensor, chi_out, right_kind)
     if right_kind == "mps":
diff --git a/src/src_method/compress.py b/src/src_method/compress.py
index b41abab..af8eacf 100644
--- a/src/src_method/compress.py
+++ b/src/src_method/compress.py
@@ -17,7 +17,12 @@
 import structlog
 from opt_einsum import contract
 
-from ._tensor_train import MIN_SRC_SITES, exact_compress, infer_kind
+from ._tensor_train import (
+    MIN_SRC_SITES,
+    check_exact_supported,
+    exact_compress,
+    infer_kind,
+)
 from .utils import (
     default_rng,
     get_xp,
@@ -87,7 +92,8 @@ def compress(
 
     Raises:
         TypeError: If the input tensor type is unsupported.
-        ValueError: If ``device`` is not recognised.
+        ValueError: If a sub-three-site train is not exactly two sites, or if
+            ``device`` is not recognised.
         ImportError: If ``device="gpu"`` but cupy is not installed.
 
     """
@@ -103,6 +109,7 @@ def compress(
         raise TypeError(msg)
 
     if len(tensor) < MIN_SRC_SITES:
+        check_exact_supported(len(tensor))
         logger.warning(LOG_WARN_SMALL)
         return exact_compress(tensor, chi_out, kind)
     if kind == "mps":
diff --git a/tests/test_package.py b/tests/test_package.py
index aa28cfc..e10fbca 100644
--- a/tests/test_package.py
+++ b/tests/test_package.py
@@ -423,6 +423,28 @@ def test_single_site_train_raises(phys_dim, chi_out, array_type):
         compress([np.ones((1, phys_dim), dtype=array_type)], chi_out=chi_out)
 
 
+def test_single_site_train_raises_without_warning(
+    phys_dim, chi_out, array_type, caplog
+):
+    """The degenerate train must raise before the fallback is announced."""
+    mpo = [np.ones((1, phys_dim, phys_dim), dtype=array_type)]
+    mps = [np.ones((1, phys_dim), dtype=array_type)]
+
+    with pytest.raises(ValueError, match="two-site tensor train"):
+        apply(mpo, mps, chi_out=chi_out)
+
+    assert "Defaulting" not in caplog.text
+
+
+def test_mismatched_site_counts_raise(phys_dim, chi_out, array_type):
+    """A shorter left train must not silently truncate the right one."""
+    H = qtn.MPO_identity(4, phys_dim=phys_dim, dtype=array_type)
+    psi = qtn.MPS_rand_state(6, bond_dim=4, phys_dim=phys_dim, dtype=array_type)
+
+    with pytest.raises(ValueError, match="same number of sites"):
+        apply(H.arrays, psi.arrays, chi_out=chi_out, dtype=array_type)
+
+
 # -------------------------------------------------
 # --- Test adaptive bond truncation via cutoff  ---
 # -------------------------------------------------