From d7608184f557118fe2d967e89a1a4c198b8fa0a4 Mon Sep 17 00:00:00 2001 From: Deshan Date: Wed, 4 Mar 2026 16:16:16 +0545 Subject: [PATCH 1/2] Add device parameter to ToCupy operator and to_cupy utility --- pylops/basicoperators/tocupy.py | 7 ++++++- pylops/utils/backend.py | 8 ++++++-- pytests/test_basicoperators.py | 16 ++++++++++++++++ 3 files changed, 28 insertions(+), 3 deletions(-) diff --git a/pylops/basicoperators/tocupy.py b/pylops/basicoperators/tocupy.py index a1dfb5b0..6d3006e2 100644 --- a/pylops/basicoperators/tocupy.py +++ b/pylops/basicoperators/tocupy.py @@ -22,6 +22,9 @@ class ToCupy(LinearOperator): Number of samples for each dimension dtype : :obj:`str`, optional Type of elements in input array. + device : :obj:`int`, optional + GPU device where the array will be transferred to. + By default, device 0 is used. name : :obj:`str`, optional Name of operator (to be used by :func:`pylops.utils.describe.describe`) @@ -52,13 +55,15 @@ def __init__( self, dims: Union[int, InputDimsLike], dtype: DTypeLike = "float64", + device: int = 0, name: str = "C", ) -> None: + self.device = device dims = _value_or_sized_to_tuple(dims) super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dims, name=name) def _matvec(self, x: NDArray) -> NDArray: - return to_cupy(x) + return to_cupy(x, device=self.device) def _rmatvec(self, x: NDArray) -> NDArray: return to_numpy(x) diff --git a/pylops/utils/backend.py b/pylops/utils/backend.py index 2cc125df..14acedb5 100644 --- a/pylops/utils/backend.py +++ b/pylops/utils/backend.py @@ -485,13 +485,16 @@ def get_real_dtype(dtype: DTypeLike) -> DTypeLike: return np.real(np.ones(1, dtype)).dtype -def to_cupy(x: ArrayLike) -> ArrayLike: +def to_cupy(x: ArrayLike, device: int = 0) -> ArrayLike: """Convert x to cupy array if cupy is available Parameters ---------- x : :obj:`numpy.ndarray`, :obj:`cupy.ndarray` or :obj:`jax.Array` Array to evaluate + device : :obj:`int`, optional + GPU device where the array will be transferred to. + by default, device 0 is used. Returns ------- @@ -501,7 +504,8 @@ def to_cupy(x: ArrayLike) -> ArrayLike: """ if deps.cupy_enabled: if cp.get_array_module(x) != cp: - x = cp.asarray(x) + with cp.cuda.Device(device): + x = cp.asarray(x) return x diff --git a/pytests/test_basicoperators.py b/pytests/test_basicoperators.py index 447cfa95..887be620 100644 --- a/pytests/test_basicoperators.py +++ b/pytests/test_basicoperators.py @@ -718,3 +718,19 @@ def test_ToCupy(par): xadj = Top.H * y assert_array_equal(x, xadj) assert_array_equal(y, np.asarray(x)) + + +@pytest.mark.parametrize("par", [(par1), (par2), (par1j), (par2j), (par3)]) +def test_ToCupy_device(par): + """Forward and adjoint for ToCupy operator with explicit device parameter + (checking device is correctly stored and used) + """ + Top = ToCupy(par["nx"], dtype=par["dtype"], device=0) + + np.random.seed(10) + x = npp.random.randn(par["nx"]) + par["imag"] * npp.random.randn(par["nx"]) + y = Top * x + xadj = Top.H * y + assert_array_equal(x, xadj) + assert_array_equal(y, np.asarray(x)) + assert Top.device == 0 From e795c51278af8f4c3e53e21e988d1fe8713e1ccb Mon Sep 17 00:00:00 2001 From: Deshan Date: Wed, 4 Mar 2026 16:31:51 +0545 Subject: [PATCH 2/2] Fix trailing whitespace in ToCupy docstring --- pylops/basicoperators/tocupy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pylops/basicoperators/tocupy.py b/pylops/basicoperators/tocupy.py index 6d3006e2..c3113973 100644 --- a/pylops/basicoperators/tocupy.py +++ b/pylops/basicoperators/tocupy.py @@ -24,7 +24,7 @@ class ToCupy(LinearOperator): Type of elements in input array. device : :obj:`int`, optional GPU device where the array will be transferred to. - By default, device 0 is used. + By default, device 0 is used. name : :obj:`str`, optional Name of operator (to be used by :func:`pylops.utils.describe.describe`)