diff --git a/src/src_method/_sweep.py b/src/src_method/_sweep.py index 3529b5d..280453b 100644 --- a/src/src_method/_sweep.py +++ b/src/src_method/_sweep.py @@ -18,7 +18,7 @@ from itertools import count from math import prod from time import perf_counter_ns -from typing import TYPE_CHECKING, NamedTuple +from typing import TYPE_CHECKING, Any, NamedTuple from opt_einsum import contract_expression, get_symbol @@ -101,7 +101,8 @@ def sweep( *, cutoff: float = 0.0, dtype: DTypeLike, -) -> list[NDArray]: + to_host: bool = True, +) -> list[Any]: """Contract and compress a stack in ket form with one SRC sweep. Args: @@ -114,10 +115,12 @@ def sweep( xp: Array module (``numpy`` or ``cupy``). cutoff: Relative singular-value cutoff for adaptive bond truncation. dtype: The data type of the sketches. + to_host: If ``True`` (default), the result is copied to the host as numpy + arrays. If ``False``, it stays on the device of ``xp``. Returns: The site arrays of the compressed train in right-canonical form, as numpy - arrays. + arrays if ``to_host`` is ``True``, otherwise as arrays of ``xp``.. """ depth = len(layers) n_sites = len(layers[0]) @@ -157,4 +160,7 @@ def sweep( eta = [first.reshape(1, *first.shape[depth:]), *reversed(eta_reversed)] logger.debug("Right-to-left sweep: %.3f s", (perf_counter_ns() - tms) * 1e-9) - return [to_numpy(site) for site in unpad(eta, kind)] + sites_out = unpad(eta, kind) + if to_host: + return [to_numpy(site) for site in sites_out] + return list(sites_out) diff --git a/src/src_method/apply.py b/src/src_method/apply.py index 534c00d..0d04f8a 100644 --- a/src/src_method/apply.py +++ b/src/src_method/apply.py @@ -8,7 +8,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from ._tensor_train import infer_kind from .stack import src @@ -28,7 +28,8 @@ def apply( dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", -) -> list[NDArray]: + to_host: bool = True, +) -> list[Any]: """Applies the Successive Randomized Compression (SRC) algorithm. Equivalent to ``src(left_tensor, right_tensor, ...)`` restricted to an MPO on @@ -54,6 +55,8 @@ def apply( seed: An optional seed for the random number generator. device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy). Requires the optional ``cupy`` dependency for GPU execution. + to_host: If ``True`` (default), results are returned as numpy arrays on + the host. Set to ``False`` to keep tensors on the compute device. Returns: The site arrays of the compressed tensor network (MPS or MPO). @@ -85,4 +88,5 @@ def apply( dtype=dtype, seed=seed, device=device, + to_host=to_host, ) diff --git a/src/src_method/compress.py b/src/src_method/compress.py index bdff465..f6f9a15 100644 --- a/src/src_method/compress.py +++ b/src/src_method/compress.py @@ -8,7 +8,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from .stack import src @@ -26,7 +26,8 @@ def compress( dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", -) -> list[NDArray]: + to_host: bool = True, +) -> list[Any]: """Applies the Successive Randomized Compression (SRC) algorithm. Equivalent to ``src(tensor, ...)``; see `src_method.stack.src` for the @@ -50,6 +51,8 @@ def compress( seed: An optional seed for the random number generator. device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy). Requires the optional ``cupy`` dependency for GPU execution. + to_host: If ``True`` (default), results are returned as numpy arrays on + the host. Set to ``False`` to keep tensors on the compute device. Returns: The site arrays of the compressed tensor network (MPS or MPO). @@ -64,5 +67,11 @@ def compress( ImportError: If ``device="gpu"`` but cupy is not installed. """ return src( - tensor, chi_out=chi_out, cutoff=cutoff, dtype=dtype, seed=seed, device=device + tensor, + chi_out=chi_out, + cutoff=cutoff, + dtype=dtype, + seed=seed, + device=device, + to_host=to_host, ) diff --git a/src/src_method/stack.py b/src/src_method/stack.py index a7f7ace..2a2b5a8 100644 --- a/src/src_method/stack.py +++ b/src/src_method/stack.py @@ -8,7 +8,7 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from ._sweep import sweep from ._tensor_train import ( @@ -40,7 +40,8 @@ def src( dtype: DTypeLike | None = None, seed: int | None = None, device: str = "cpu", -) -> list[NDArray]: + to_host: bool = True, +) -> list[Any]: """Contract a stack of tensor trains and compress the result with SRC. The stack ``src(T_1, T_2, ..., T_m)`` is the product ``T_1 T_2 ... T_m`` in @@ -80,10 +81,12 @@ def src( seed: An optional seed for the random number generator. device: ``"cpu"`` (default, numpy) or ``"gpu"`` (cupy). Requires the optional ``cupy`` dependency for GPU execution. + to_host: If ``True`` (default), results are returned as numpy arrays on + the host. Set to ``False`` to keep tensors on the compute device. Returns: The site arrays of the compressed train (MPS or MPO), in right-canonical - form, as numpy arrays (host-side, whatever the ``device``). + form, as numpy arrays if ``to_host`` is ``True``, otherwise as arrays of ``xp``. Raises: TypeError: If ``chi_out`` is not an integer, if a train has an unrecognised @@ -105,7 +108,8 @@ def src( if n_sites < MIN_SRC_SITES: check_exact_supported(n_sites) logger.warning(LOG_WARN_SMALL) - return exact_stack(layers, chi_out, kind) + result = exact_stack(layers, chi_out, kind) + return result if to_host else [xp.asarray(t) for t in result] logger.debug( "Starting SRC: n_sites=%d, depth=%d, output=%s, device=%s", @@ -122,6 +126,7 @@ def src( xp, cutoff=cutoff, dtype=sketch_dtype(dtype, *layers), + to_host=to_host, ) logger.debug("SRC complete") return result diff --git a/src/src_method/utils/linalg.py b/src/src_method/utils/linalg.py index ff2e89c..bbc5935 100644 --- a/src/src_method/utils/linalg.py +++ b/src/src_method/utils/linalg.py @@ -40,16 +40,13 @@ def truncated_qr( return xp.linalg.qr(matrix)[0] Q, R = xp.linalg.qr(matrix.T if transpose else matrix) - R_np = R.get() if hasattr(R, "get") else np.asarray(R) - U, S, _ = np.linalg.svd(R_np.T if transpose else R_np, full_matrices=False) + U, S, _ = xp.linalg.svd(R.T if transpose else R, full_matrices=False) rank = max(1, int((cutoff * S[0] <= S).sum())) if transpose: Q_trunc = U[:, :rank] else: - U_trunc = xp.asarray(U[:, :rank]) if xp is not np else U[:, :rank] + U_trunc = U[:, :rank] Q_trunc = Q @ U_trunc - if xp is not np: - return xp.asarray(Q_trunc) return Q_trunc diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index 420604a..cb10c80 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -181,3 +181,115 @@ def test_src_stack_gpu_matches_reference(device: str) -> None: ref = H1.apply(H2.apply(psi, compress=False), compress=False) np.testing.assert_allclose(ref.distance(out), 0.0, atol=1e-6) + + +def test_apply_mpo_mps_to_host_false_returns_cupy() -> None: + """With to_host=False, apply() must return cupy arrays on GPU.""" + n_sites, phys_dim, chi_out = 5, 2, 8 + H = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128) + psi = qtn.MPS_rand_state( + n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128 + ) + + result = apply( + H.arrays, + psi.arrays, + chi_out=chi_out, + dtype=np.complex128, + seed=0, + device="gpu", + to_host=False, + ) + assert all(isinstance(t, cupy.ndarray) for t in result) + + +def test_compress_mpo_to_host_false_returns_cupy() -> None: + """With to_host=False, compress() must return cupy arrays on GPU.""" + n_sites, phys_dim, chi_out = 5, 2, 8 + H = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128) + + result = compress( + H.arrays, + chi_out=chi_out, + dtype=np.complex128, + seed=0, + device="gpu", + to_host=False, + ) + assert all(isinstance(t, cupy.ndarray) for t in result) + + +def test_chained_apply_no_host_roundtrip() -> None: + """Chained GPU calls with to_host=False must work and match CPU results.""" + n_sites, phys_dim, chi_out = 5, 2, 8 + H1 = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128) + H2 = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128) + psi = qtn.MPS_rand_state( + n_sites, bond_dim=chi_out, phys_dim=phys_dim, dtype=np.complex128 + ) + + # 1. Chained GPU apply + intermediate_gpu = apply( + H1.arrays, + psi.arrays, + chi_out=chi_out, + dtype=np.complex128, + seed=0, + device="gpu", + to_host=False, + ) + assert isinstance(intermediate_gpu[0], cupy.ndarray) + + final_gpu = apply( + H2.arrays, + intermediate_gpu, + chi_out=chi_out, + dtype=np.complex128, + seed=0, + device="gpu", + to_host=False, + ) + assert isinstance(final_gpu[0], cupy.ndarray) + + # 2. Chained CPU apply for numerical comparison + intermediate_cpu = apply( + H1.arrays, + psi.arrays, + chi_out=chi_out, + dtype=np.complex128, + seed=0, + device="cpu", + to_host=True, + ) + final_cpu = apply( + H2.arrays, + intermediate_cpu, + chi_out=chi_out, + dtype=np.complex128, + seed=0, + device="cpu", + to_host=True, + ) + + # Bring GPU result to host and compare + final_gpu_host = [t.get() for t in final_gpu] + + qtn_cpu = as_mps(final_cpu) + qtn_gpu = as_mps(final_gpu_host) + np.testing.assert_allclose(qtn_cpu.distance(qtn_gpu), 0.0, atol=1e-6) + + +def test_to_host_true_backward_compatible() -> None: + """to_host=True (implicit default) must return numpy arrays even when device='gpu'.""" + n_sites, phys_dim, chi_out = 5, 2, 8 + H = qtn.MPO_rand(n_sites, bond_dim=4, phys_dim=phys_dim, dtype=np.complex128) + + result = compress( + H.arrays, + chi_out=chi_out, + dtype=np.complex128, + seed=0, + device="gpu", + # to_host=True is the implicit default here + ) + assert all(isinstance(t, np.ndarray) for t in result)