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..b16baec --- /dev/null +++ b/src/cunumpy/interop/_common.py @@ -0,0 +1,96 @@ +"""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. + + 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: + 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. + + Parameters + ---------- + source, target : object + Source and converted arrays. + """ + 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. + + Parameters + ---------- + backend : {"numpy", "cupy"} + Name of the target array library. + + Returns + ------- + module + The requested array 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. + + 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 new file mode 100644 index 0000000..4258974 --- /dev/null +++ b/src/cunumpy/interop/dlpack.py @@ -0,0 +1,54 @@ +"""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`. + + Parameters + ---------- + array : object + Array implementing the DLPack protocol. + + Returns + ------- + object + The DLPack capsule or protocol object. + """ + 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. + + 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") + 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..98393f4 --- /dev/null +++ b/src/cunumpy/interop/jax.py @@ -0,0 +1,76 @@ +"""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. + + 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 + + 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. + + 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) + 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..976978c --- /dev/null +++ b/src/cunumpy/interop/mlx.py @@ -0,0 +1,71 @@ +"""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. + + 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 + 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. + + 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) + 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..6b77123 --- /dev/null +++ b/src/cunumpy/interop/torch.py @@ -0,0 +1,72 @@ +"""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. + + 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 + 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. + + 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) + 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..22442ff --- /dev/null +++ b/tests/unit/test_interop.py @@ -0,0 +1,39 @@ +"""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")