From 4dbb1c28e695387cc92fb588cabee0d8aab6d2af Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 22:41:51 +0200 Subject: [PATCH 1/3] Implemented optional interop with jax, torch mlx --- pyproject.toml | 3 ++ src/cunumpy/__init__.py | 2 ++ src/cunumpy/__init__.pyi | 1 + src/cunumpy/interop/__init__.py | 9 ++++++ src/cunumpy/interop/_common.py | 53 +++++++++++++++++++++++++++++++++ src/cunumpy/interop/dlpack.py | 30 +++++++++++++++++++ src/cunumpy/interop/jax.py | 49 ++++++++++++++++++++++++++++++ src/cunumpy/interop/mlx.py | 43 ++++++++++++++++++++++++++ src/cunumpy/interop/torch.py | 44 +++++++++++++++++++++++++++ tests/unit/test_interop.py | 41 +++++++++++++++++++++++++ 10 files changed, 275 insertions(+) create mode 100644 src/cunumpy/interop/__init__.py create mode 100644 src/cunumpy/interop/_common.py create mode 100644 src/cunumpy/interop/dlpack.py create mode 100644 src/cunumpy/interop/jax.py create mode 100644 src/cunumpy/interop/mlx.py create mode 100644 src/cunumpy/interop/torch.py create mode 100644 tests/unit/test_interop.py diff --git a/pyproject.toml b/pyproject.toml index f73a9f1..dda36b0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,9 @@ optional-dependencies.dev = [ # the shared docs workflow registers an ipykernel even without notebooks optional-dependencies.docs = [ "griffe", "ipykernel" ] optional-dependencies.metal = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ] +optional-dependencies.jax = [ "jax" ] +optional-dependencies.mlx = [ "mlx; sys_platform == 'darwin' and platform_machine == 'arm64'" ] +optional-dependencies.torch = [ "torch" ] optional-dependencies.test = [ "coverage", "pytest", "scipy" ] optional-dependencies.test-compiled = [ "cunumpy[test]", "pyccel" ] urls."Source" = "https://github.com/max-models/cunumpy" diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index d8743db..d240d97 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -22,6 +22,7 @@ arguments, cuda, integrate, + interop, kernels, memory, mpi, @@ -123,6 +124,7 @@ def require_version(minimum: str) -> None: "get_backend", "host_call", "integrate", + "interop", "is_cpu", "is_gpu", "kernels", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index e239928..e89ee4c 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -12,6 +12,7 @@ from cunumpy import algorithms as algorithms from cunumpy import arguments as arguments from cunumpy import cuda as cuda from cunumpy import integrate as integrate +from cunumpy import interop as interop from cunumpy import kernels as kernels from cunumpy import memory as memory from cunumpy import mpi as mpi diff --git a/src/cunumpy/interop/__init__.py b/src/cunumpy/interop/__init__.py new file mode 100644 index 0000000..e7deedb --- /dev/null +++ b/src/cunumpy/interop/__init__.py @@ -0,0 +1,9 @@ +"""Explicit interoperability with scientific machine-learning libraries. + +The adapters are optional and import their framework lazily. They do not +participate in CuNumpy backend selection. +""" + +from cunumpy.interop import dlpack, jax, mlx, torch + +__all__ = ["dlpack", "jax", "mlx", "torch"] diff --git a/src/cunumpy/interop/_common.py b/src/cunumpy/interop/_common.py new file mode 100644 index 0000000..05bf591 --- /dev/null +++ b/src/cunumpy/interop/_common.py @@ -0,0 +1,53 @@ +"""Shared implementation for optional framework adapters.""" + +from __future__ import annotations + +from importlib import import_module +from typing import Any + +import numpy as np + +from cunumpy._transfers import _ACTIVE, _describe, _nbytes, _record + + +def require(module: str, *, extra: str | None = None) -> Any: + """Import an optional framework and give a consistent installation error.""" + try: + return import_module(module) + except ImportError as error: + extra = module.split(".", 1)[0] if extra is None else extra + raise ImportError( + f"{module} interoperability needs {module.split('.', 1)[0]!r}; " + f"install cunumpy[{extra}]" + ) from error + + +def record(source: Any, target: Any) -> None: + """Record an adapter conversion when transfer profiling is active.""" + if _ACTIVE: + _record( + "device_copy", + f"interop {type(source).__name__} -> {type(target).__name__} " + f"({_describe(target)})", + nbytes=_nbytes(target), + ) + + +def target_backend(backend: str) -> Any: + """Load a NumPy or CuPy target module.""" + if backend == "numpy": + return np + if backend == "cupy": + try: + import cupy + except ImportError as error: + raise ImportError( + "CuPy interoperability needs CuPy; install a compatible CuPy release" + ) from error + return cupy + raise ValueError(f"backend must be 'numpy' or 'cupy', got {backend!r}") + + +def cast(array: Any, dtype: Any) -> Any: + """Cast an array only when a dtype was requested.""" + return array if dtype is None else array.astype(dtype) diff --git a/src/cunumpy/interop/dlpack.py b/src/cunumpy/interop/dlpack.py new file mode 100644 index 0000000..bed9333 --- /dev/null +++ b/src/cunumpy/interop/dlpack.py @@ -0,0 +1,30 @@ +"""Small explicit helpers around the DLPack array interchange protocol.""" + +from __future__ import annotations + +from typing import Any + +from cunumpy.interop._common import record, target_backend + + +def to_dlpack(array: Any) -> Any: + """Return a DLPack capsule or protocol object for `array`.""" + try: + return array.__dlpack__() + except AttributeError as error: + raise TypeError( + f"{type(array).__name__} does not implement the DLPack protocol" + ) from error + + +def from_dlpack(source: Any, *, backend: str = "numpy") -> Any: + """Construct a NumPy or CuPy array from a DLPack source.""" + target = target_backend(backend) + if not hasattr(target, "from_dlpack"): + raise TypeError(f"{backend} does not provide from_dlpack") + result = target.from_dlpack(source) + record(source, result) + return result + + +__all__ = ["from_dlpack", "to_dlpack"] diff --git a/src/cunumpy/interop/jax.py b/src/cunumpy/interop/jax.py new file mode 100644 index 0000000..af48fea --- /dev/null +++ b/src/cunumpy/interop/jax.py @@ -0,0 +1,49 @@ +"""Optional interoperability with JAX arrays.""" + +from __future__ import annotations + +from typing import Any + +from cunumpy.interop._common import cast, record, require, target_backend + + +def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: + """Convert a NumPy or CuPy array to a JAX array. + + DLPack is used when available; `device` and `dtype` are explicit optional + conversions. + """ + jax = require("jax") + import jax.numpy as jnp + + if type(array).__module__.startswith("jax"): + result = array + elif hasattr(array, "__dlpack__"): + result = jax.dlpack.from_dlpack(array) + else: + result = jnp.asarray(array) + result = cast(result, dtype) + if device is not None: + result = jax.device_put(result, device) + record(array, result) + return result + + +def from_framework( + array: Any, *, backend: str = "numpy", dtype: Any = None +) -> Any: + """Convert a JAX array to a NumPy or CuPy array.""" + target = target_backend(backend) + if hasattr(target, "from_dlpack") and hasattr(array, "__dlpack__"): + result = target.from_dlpack(array) + else: + result = target.asarray(array) + result = cast(result, dtype) + record(array, result) + return result + + +to_jax = to_framework +from_jax = from_framework + +__all__ = ["from_framework", "from_jax", "to_framework", "to_jax"] diff --git a/src/cunumpy/interop/mlx.py b/src/cunumpy/interop/mlx.py new file mode 100644 index 0000000..e3b347b --- /dev/null +++ b/src/cunumpy/interop/mlx.py @@ -0,0 +1,43 @@ +"""Optional interoperability with Apple's MLX arrays.""" + +from __future__ import annotations + +from typing import Any + +from cunumpy.interop._common import cast, record, require, target_backend + + +def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: + """Convert a NumPy or CuPy array to an MLX array.""" + mlx = require("mlx.core", extra="mlx") + if type(array).__module__.startswith("mlx"): + result = array + elif hasattr(mlx, "from_dlpack") and hasattr(array, "__dlpack__"): + result = mlx.from_dlpack(array) + else: + result = mlx.array(array) + result = cast(result, dtype) + if device is not None: + result = mlx.array(result, device=device) + record(array, result) + return result + + +def from_framework( + array: Any, *, backend: str = "numpy", dtype: Any = None +) -> Any: + """Convert an MLX array to a NumPy or CuPy array.""" + target = target_backend(backend) + if hasattr(array, "__dlpack__") and hasattr(target, "from_dlpack"): + result = target.from_dlpack(array) + else: + result = target.asarray(array) + result = cast(result, dtype) + record(array, result) + return result + + +to_mlx = to_framework +from_mlx = from_framework + +__all__ = ["from_framework", "from_mlx", "to_framework", "to_mlx"] diff --git a/src/cunumpy/interop/torch.py b/src/cunumpy/interop/torch.py new file mode 100644 index 0000000..27791d6 --- /dev/null +++ b/src/cunumpy/interop/torch.py @@ -0,0 +1,44 @@ +"""Optional interoperability with PyTorch tensors.""" + +from __future__ import annotations + +from typing import Any + +from cunumpy.interop._common import cast, record, require, target_backend + + +def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: + """Convert a NumPy or CuPy array to a PyTorch tensor.""" + torch = require("torch") + if type(array).__module__.startswith("torch"): + result = array + elif hasattr(torch, "from_dlpack") and hasattr(array, "__dlpack__"): + result = torch.from_dlpack(array) + else: + result = torch.as_tensor(array) + if dtype is not None: + result = result.to(dtype=dtype) + if device is not None: + result = result.to(device=device) + record(array, result) + return result + + +def from_framework( + array: Any, *, backend: str = "numpy", dtype: Any = None +) -> Any: + """Convert a PyTorch tensor to a NumPy or CuPy array.""" + target = target_backend(backend) + if hasattr(target, "from_dlpack") and hasattr(array, "__dlpack__"): + result = target.from_dlpack(array) + else: + result = target.asarray(array.detach().cpu().numpy()) + result = cast(result, dtype) + record(array, result) + return result + + +to_torch = to_framework +from_torch = from_framework + +__all__ = ["from_framework", "from_torch", "to_framework", "to_torch"] diff --git a/tests/unit/test_interop.py b/tests/unit/test_interop.py new file mode 100644 index 0000000..185e4d2 --- /dev/null +++ b/tests/unit/test_interop.py @@ -0,0 +1,41 @@ +"""Tests for optional framework interoperability.""" + +import types + +import numpy as np +import pytest + +from cunumpy.interop import dlpack +from cunumpy.interop._common import require, target_backend + + +def test_dlpack_helpers_use_array_protocol(monkeypatch): + class Source: + def __dlpack__(self): + return "capsule" + + assert dlpack.to_dlpack(Source()) == "capsule" + + target = types.SimpleNamespace( + from_dlpack=lambda source: np.array([source.__dlpack__()]), + ) + monkeypatch.setattr( + "cunumpy.interop.dlpack.target_backend", lambda backend: target + ) + result = dlpack.from_dlpack(Source()) + np.testing.assert_array_equal(result, ["capsule"]) + + +def test_dlpack_rejects_objects_without_protocol(): + with pytest.raises(TypeError, match="DLPack protocol"): + dlpack.to_dlpack(object()) + + +def test_optional_dependency_error_is_actionable(): + with pytest.raises(ImportError, match="cunumpy\\[not-installed-framework\\]"): + require("not-installed-framework") + + +def test_target_backend_rejects_unknown_backend(): + with pytest.raises(ValueError, match="backend must be"): + target_backend("jax") From bcc9c1917688d76c95365c856e3fc2506fd633b7 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 22:46:52 +0200 Subject: [PATCH 2/3] ruff format --- src/cunumpy/interop/jax.py | 4 +--- src/cunumpy/interop/mlx.py | 4 +--- src/cunumpy/interop/torch.py | 4 +--- tests/unit/test_interop.py | 4 +--- 4 files changed, 4 insertions(+), 12 deletions(-) diff --git a/src/cunumpy/interop/jax.py b/src/cunumpy/interop/jax.py index af48fea..09e416f 100644 --- a/src/cunumpy/interop/jax.py +++ b/src/cunumpy/interop/jax.py @@ -29,9 +29,7 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: return result -def from_framework( - array: Any, *, backend: str = "numpy", dtype: Any = None -) -> Any: +def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: """Convert a JAX array to a NumPy or CuPy array.""" target = target_backend(backend) if hasattr(target, "from_dlpack") and hasattr(array, "__dlpack__"): diff --git a/src/cunumpy/interop/mlx.py b/src/cunumpy/interop/mlx.py index e3b347b..e62ee9d 100644 --- a/src/cunumpy/interop/mlx.py +++ b/src/cunumpy/interop/mlx.py @@ -23,9 +23,7 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: return result -def from_framework( - array: Any, *, backend: str = "numpy", dtype: Any = None -) -> Any: +def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: """Convert an MLX array to a NumPy or CuPy array.""" target = target_backend(backend) if hasattr(array, "__dlpack__") and hasattr(target, "from_dlpack"): diff --git a/src/cunumpy/interop/torch.py b/src/cunumpy/interop/torch.py index 27791d6..87ee526 100644 --- a/src/cunumpy/interop/torch.py +++ b/src/cunumpy/interop/torch.py @@ -24,9 +24,7 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: return result -def from_framework( - array: Any, *, backend: str = "numpy", dtype: Any = None -) -> Any: +def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: """Convert a PyTorch tensor to a NumPy or CuPy array.""" target = target_backend(backend) if hasattr(target, "from_dlpack") and hasattr(array, "__dlpack__"): diff --git a/tests/unit/test_interop.py b/tests/unit/test_interop.py index 185e4d2..22442ff 100644 --- a/tests/unit/test_interop.py +++ b/tests/unit/test_interop.py @@ -19,9 +19,7 @@ def __dlpack__(self): target = types.SimpleNamespace( from_dlpack=lambda source: np.array([source.__dlpack__()]), ) - monkeypatch.setattr( - "cunumpy.interop.dlpack.target_backend", lambda backend: target - ) + monkeypatch.setattr("cunumpy.interop.dlpack.target_backend", lambda backend: target) result = dlpack.from_dlpack(Source()) np.testing.assert_array_equal(result, ["capsule"]) From ec989320ea5cbdb3b33057741ce172fef0f15f34 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 22:57:44 +0200 Subject: [PATCH 3/3] numpydoc --- src/cunumpy/interop/_common.py | 51 +++++++++++++++++++++++++++++++--- src/cunumpy/interop/dlpack.py | 28 +++++++++++++++++-- src/cunumpy/interop/jax.py | 31 ++++++++++++++++++++- src/cunumpy/interop/mlx.py | 34 +++++++++++++++++++++-- src/cunumpy/interop/torch.py | 34 +++++++++++++++++++++-- 5 files changed, 167 insertions(+), 11 deletions(-) diff --git a/src/cunumpy/interop/_common.py b/src/cunumpy/interop/_common.py index 05bf591..b16baec 100644 --- a/src/cunumpy/interop/_common.py +++ b/src/cunumpy/interop/_common.py @@ -11,7 +11,20 @@ def require(module: str, *, extra: str | None = None) -> Any: - """Import an optional framework and give a consistent installation error.""" + """Import an optional framework and give a consistent installation error. + + Parameters + ---------- + module : str + Import name of the optional framework. + extra : str or None, optional + CuNumpy extra name, inferred from `module` when omitted. + + Returns + ------- + module + Imported framework module. + """ try: return import_module(module) except ImportError as error: @@ -23,7 +36,13 @@ def require(module: str, *, extra: str | None = None) -> Any: def record(source: Any, target: Any) -> None: - """Record an adapter conversion when transfer profiling is active.""" + """Record an adapter conversion when transfer profiling is active. + + Parameters + ---------- + source, target : object + Source and converted arrays. + """ if _ACTIVE: _record( "device_copy", @@ -34,7 +53,18 @@ def record(source: Any, target: Any) -> None: def target_backend(backend: str) -> Any: - """Load a NumPy or CuPy target module.""" + """Load a NumPy or CuPy target module. + + Parameters + ---------- + backend : {"numpy", "cupy"} + Name of the target array library. + + Returns + ------- + module + The requested array module. + """ if backend == "numpy": return np if backend == "cupy": @@ -49,5 +79,18 @@ def target_backend(backend: str) -> Any: def cast(array: Any, dtype: Any) -> Any: - """Cast an array only when a dtype was requested.""" + """Cast an array only when a dtype was requested. + + Parameters + ---------- + array : object + Array to cast. + dtype : object or None + Requested dtype, or None to preserve the input. + + Returns + ------- + object + The original or cast array. + """ return array if dtype is None else array.astype(dtype) diff --git a/src/cunumpy/interop/dlpack.py b/src/cunumpy/interop/dlpack.py index bed9333..4258974 100644 --- a/src/cunumpy/interop/dlpack.py +++ b/src/cunumpy/interop/dlpack.py @@ -8,7 +8,18 @@ def to_dlpack(array: Any) -> Any: - """Return a DLPack capsule or protocol object for `array`.""" + """Return a DLPack capsule or protocol object for `array`. + + Parameters + ---------- + array : object + Array implementing the DLPack protocol. + + Returns + ------- + object + The DLPack capsule or protocol object. + """ try: return array.__dlpack__() except AttributeError as error: @@ -18,7 +29,20 @@ def to_dlpack(array: Any) -> Any: def from_dlpack(source: Any, *, backend: str = "numpy") -> Any: - """Construct a NumPy or CuPy array from a DLPack source.""" + """Construct a NumPy or CuPy array from a DLPack source. + + Parameters + ---------- + source : object + DLPack capsule or object implementing the protocol. + backend : {"numpy", "cupy"}, default="numpy" + Array library to construct. + + Returns + ------- + object + The constructed array. + """ target = target_backend(backend) if not hasattr(target, "from_dlpack"): raise TypeError(f"{backend} does not provide from_dlpack") diff --git a/src/cunumpy/interop/jax.py b/src/cunumpy/interop/jax.py index 09e416f..98393f4 100644 --- a/src/cunumpy/interop/jax.py +++ b/src/cunumpy/interop/jax.py @@ -12,6 +12,20 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: DLPack is used when available; `device` and `dtype` are explicit optional conversions. + + Parameters + ---------- + array : object + Array to convert. + device : object or None, optional + JAX device on which to place the result. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + JAX array. """ jax = require("jax") import jax.numpy as jnp @@ -30,7 +44,22 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: - """Convert a JAX array to a NumPy or CuPy array.""" + """Convert a JAX array to a NumPy or CuPy array. + + Parameters + ---------- + array : object + JAX array to convert. + backend : {"numpy", "cupy"}, default="numpy" + Array library to construct. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + NumPy or CuPy array. + """ target = target_backend(backend) if hasattr(target, "from_dlpack") and hasattr(array, "__dlpack__"): result = target.from_dlpack(array) diff --git a/src/cunumpy/interop/mlx.py b/src/cunumpy/interop/mlx.py index e62ee9d..976978c 100644 --- a/src/cunumpy/interop/mlx.py +++ b/src/cunumpy/interop/mlx.py @@ -8,7 +8,22 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: - """Convert a NumPy or CuPy array to an MLX array.""" + """Convert a NumPy or CuPy array to an MLX array. + + Parameters + ---------- + array : object + Array to convert. + device : object or None, optional + MLX device on which to place the result. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + MLX array. + """ mlx = require("mlx.core", extra="mlx") if type(array).__module__.startswith("mlx"): result = array @@ -24,7 +39,22 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: - """Convert an MLX array to a NumPy or CuPy array.""" + """Convert an MLX array to a NumPy or CuPy array. + + Parameters + ---------- + array : object + MLX array to convert. + backend : {"numpy", "cupy"}, default="numpy" + Array library to construct. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + NumPy or CuPy array. + """ target = target_backend(backend) if hasattr(array, "__dlpack__") and hasattr(target, "from_dlpack"): result = target.from_dlpack(array) diff --git a/src/cunumpy/interop/torch.py b/src/cunumpy/interop/torch.py index 87ee526..6b77123 100644 --- a/src/cunumpy/interop/torch.py +++ b/src/cunumpy/interop/torch.py @@ -8,7 +8,22 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: - """Convert a NumPy or CuPy array to a PyTorch tensor.""" + """Convert a NumPy or CuPy array to a PyTorch tensor. + + Parameters + ---------- + array : object + Array to convert. + device : object or None, optional + PyTorch device on which to place the result. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + PyTorch tensor. + """ torch = require("torch") if type(array).__module__.startswith("torch"): result = array @@ -25,7 +40,22 @@ def to_framework(array: Any, *, device: Any = None, dtype: Any = None) -> Any: def from_framework(array: Any, *, backend: str = "numpy", dtype: Any = None) -> Any: - """Convert a PyTorch tensor to a NumPy or CuPy array.""" + """Convert a PyTorch tensor to a NumPy or CuPy array. + + Parameters + ---------- + array : object + PyTorch tensor to convert. + backend : {"numpy", "cupy"}, default="numpy" + Array library to construct. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + NumPy or CuPy array. + """ target = target_backend(backend) if hasattr(target, "from_dlpack") and hasattr(array, "__dlpack__"): result = target.from_dlpack(array)