Fix/14 gpu round trip

src/src_method/_sweep.py

@@ -19,7 +19,7 @@
 
 import logging
 from time import perf_counter_ns
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -64,8 +64,9 @@ def sweep(
     *,
     cutoff: float = 0.0,
     dtype: DTypeLike,
+    to_host: bool = True,
     resources: Resources | None = None,
-) -> list[NDArray]:
+) -> list[Any]:
     """Contract and compress a stack in ket form with one SRC sweep.
 
     Args:
@@ -80,11 +81,13 @@ 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``.
         resources: Memory budgets and scratch space; detected when ``None``.
 
     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``.
     """
     shapes = padded_shapes(layers)
     # The dtype of the environments and the output: the sketches promoted by the
@@ -140,7 +143,12 @@ def sweep(
             store.stall_seconds,
         )
         logger.debug("Device pool: %d B", device_pool_bytes(xp))
-    return [to_numpy(site) for site in unpad(eta, kind)]
+
+    # La tua fix per il GPU round-trip applicata al nuovo codice:
+    sites_out = unpad(eta, kind)
+    if to_host:
+        return [to_numpy(site) for site in sites_out]
+    return list(sites_out)
 
 
 def _left_to_right(
@@ -211,10 +219,10 @@ def _right_to_left(
         for lo, hi in _ranges(Q.shape[1], site.project_batch):
             new_proj[lo:hi] = kernels.project(eta_j[lo:hi].conj(), cores, proj)
         proj = new_proj
-        eta_reversed.append(to_numpy(eta_j))
+        eta_reversed.append(eta_j)
     cores = source[0]
     up, down = cores[0].shape[2], cores[-1].shape[3]
     first = xp.empty((proj.shape[0], up, down), dtype=dtype)
     for lo, hi in _ranges(proj.shape[0], plan.sites[0].project_batch):
         first[lo:hi] = kernels.first(cores, proj[lo:hi]).reshape(hi - lo, up, down)
-    return [to_numpy(first)[None], *reversed(eta_reversed)]
+    return [first[None], *reversed(eta_reversed)]

src/src_method/_tensor_train.py

@@ -22,10 +22,11 @@
 from opt_einsum import contract
 
 from ._validation import validate_open_boundary
-from .utils import to_numpy
+from .utils import is_host, to_numpy
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
+    from types import ModuleType
 
     from numpy.typing import NDArray
 
@@ -153,8 +154,18 @@ def check_exact_supported(n_sites: int) -> None:
         raise ValueError(msg)
 
 
+def _to_backend(site: Site, xp: ModuleType) -> NDArray:
+    """Ensure a site tensor is an array on the backend ``xp``."""
+    if is_host(xp):
+        return to_numpy(site)
+    return site if isinstance(site, xp.ndarray) else xp.asarray(site)
+
+
 def exact_compress(
-    arrays: Sequence[NDArray], chi_out: int, kind: TrainKind
+    arrays: Sequence[Site],
+    chi_out: int,
+    kind: TrainKind,
+    xp: ModuleType = np,
 ) -> list[NDArray]:
     """Compress a two-site train exactly via a single truncated SVD.
 
@@ -164,32 +175,36 @@ def exact_compress(
         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"``.
+        xp: Array module (``numpy`` or ``cupy``); defaults to numpy.
 
     Returns:
-        The compressed train, in right-canonical form, as numpy arrays.
+        The compressed train, in right-canonical form, as arrays of ``xp``.
     """
-    # The dense SVD is host-side, so accept device arrays like the sweep does.
-    arrays = [to_numpy(arr) for arr in arrays]
+    device_arrays = [_to_backend(arr, xp) for arr in 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)
+        theta = contract("ab,ac->bc", device_arrays[0], device_arrays[1])
+        left, right = _truncated_svd(theta, chi_out, xp=xp)
         return [left.T, right]
 
     # (b, u0, d0) x (b, u1, d1) -> (u0, d0, u1, d1)
-    theta = contract("aij,akl->ijkl", arrays[0], arrays[1])
+    theta = contract("aij,akl->ijkl", device_arrays[0], device_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)
+    left, right = _truncated_svd(
+        theta.reshape(up_l * down_l, up_r * down_r), chi_out, xp=xp
+    )
     rank = left.shape[1]
     return [
         left.reshape(up_l, down_l, rank).transpose(2, 0, 1),
         right.reshape(rank, up_r, down_r),
     ]
 
 
-def _truncated_svd(theta: NDArray, chi_out: int) -> tuple[NDArray, NDArray]:
+def _truncated_svd(
+    theta: NDArray, chi_out: int, xp: ModuleType = np
+) -> 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)
+    U, S, Vh = xp.linalg.svd(theta, full_matrices=False)
     rank = min(chi_out, S.size)
     return U[:, :rank] * S[:rank], Vh[:rank]
 
@@ -397,7 +412,10 @@ def _check_physical_dims(
 
 
 def exact_stack(
-    layers: Sequence[Sequence[Site]], chi_out: int, kind: TrainKind
+    layers: Sequence[Sequence[Site]],
+    chi_out: int,
+    kind: TrainKind,
+    xp: ModuleType = np,
 ) -> list[NDArray]:
     """Contract and compress a two-site stack exactly.
 
@@ -409,15 +427,16 @@ def exact_stack(
         layers: A stack in ket form, as returned by `normalize_stack`.
         chi_out: The maximum bond dimension to keep.
         kind: The kind of the contracted train.
+        xp: Array module (``numpy`` or ``cupy``); defaults to numpy.
 
     Returns:
-        The compressed product, in right-canonical form, as numpy arrays.
+        The compressed product, in right-canonical form, as arrays of ``xp``.
     """
-    padded = [pad([to_numpy(site) for site in layer]) for layer in layers]
+    padded = [pad([_to_backend(site, xp) for site in layer]) for layer in layers]
     product = padded[-1]
     for layer in reversed(padded[:-1]):
         product = [_fuse(upper, lower) for upper, lower in zip(layer, product)]
-    return exact_compress(unpad(product, kind), chi_out, kind)
+    return exact_compress(unpad(product, kind), chi_out, kind, xp=xp)
 
 
 def _fuse(upper: NDArray, lower: NDArray) -> NDArray:

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
@@ -17,7 +17,7 @@
     from collections.abc import Sequence
 
     import numpy as np
-    from numpy.typing import DTypeLike, NDArray
+    from numpy.typing import DTypeLike
 
     from ._plan import Resources
     from ._tensor_train import Site
@@ -32,8 +32,9 @@ def apply(
     dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
+    to_host: bool = True,
     resources: Resources | None = None,
-) -> list[NDArray]:
+) -> list[Any]:
     """Applies the Successive Randomized Compression (SRC) algorithm.
 
     Equivalent to ``src(left_tensor, right_tensor, ...)`` restricted to an MPO on
@@ -59,6 +60,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.
         resources: Memory budgets and scratch space; see `src_method.stack.src`.
 
     Returns:
@@ -91,5 +94,6 @@ def apply(
         dtype=dtype,
         seed=seed,
         device=device,
+        to_host=to_host,
         resources=resources,
     )

src/src_method/compress.py

@@ -8,15 +8,15 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from .stack import src
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
     import numpy as np
-    from numpy.typing import DTypeLike, NDArray
+    from numpy.typing import DTypeLike
 
     from ._plan import Resources
     from ._tensor_train import Site
@@ -30,8 +30,9 @@ def compress(
     dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
+    to_host: bool = True,
     resources: Resources | None = None,
-) -> list[NDArray]:
+) -> list[Any]:
     """Applies the Successive Randomized Compression (SRC) algorithm.
 
     Equivalent to ``src(tensor, ...)``; see `src_method.stack.src` for the
@@ -55,6 +56,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.
         resources: Memory budgets and scratch space; see `src_method.stack.src`.
 
     Returns:
@@ -76,5 +79,6 @@ def compress(
         dtype=dtype,
         seed=seed,
         device=device,
+        to_host=to_host,
         resources=resources,
     )

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 (
@@ -18,13 +18,13 @@
     normalize_stack,
 )
 from ._validation import validate_chi_out, validate_cutoff
-from .utils import default_rng, get_xp, sketch_dtype
+from .utils import default_rng, get_xp, sketch_dtype, to_numpy
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
     import numpy as np
-    from numpy.typing import DTypeLike, NDArray
+    from numpy.typing import DTypeLike
 
     from ._plan import Resources
     from ._tensor_train import Site
@@ -44,8 +44,9 @@ def src(
     dtype: DTypeLike | None = None,
     seed: int | None = None,
     device: str = "cpu",
+    to_host: bool = True,
     resources: Resources | None = None,
-) -> list[NDArray]:
+) -> 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
@@ -88,13 +89,15 @@ 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.
         resources: Memory budgets and scratch space for the sweep (see
             `Resources`); every budget left unset is detected. Ignored for
             two-site stacks.
 
     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
@@ -116,7 +119,10 @@ 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, xp=xp)
+        if to_host:
+            return [to_numpy(site) for site in result]
+        return list(result)
 
     logger.debug(
         "Starting SRC: n_sites=%d, depth=%d, output=%s, device=%s",
@@ -133,6 +139,7 @@ def src(
         xp,
         cutoff=cutoff,
         dtype=sketch_dtype(dtype, *layers),
+        to_host=to_host,
         resources=resources,
     )
     logger.debug("SRC complete")

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