Implemented optional interop with jax, torch mlx

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"

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",

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

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"]

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)

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"]

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"]

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"]