From a3db37e84103da9d0b3de1f95010b509f65901c1 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 31 Oct 2026 23:02:07 +0100 Subject: [PATCH] Convert nested containers and scalar objectives in the interop adapters. --- docs/astro.config.mjs | 1 + docs/src/content/docs/guides/interop.md | 43 +++++++++ src/cunumpy/interop/_common.py | 79 ++++++++++++++++ src/cunumpy/interop/jax.py | 100 ++++++++++++++++++++ src/cunumpy/interop/mlx.py | 92 +++++++++++++++++++ src/cunumpy/interop/torch.py | 94 +++++++++++++++++++ tests/unit/test_interop_trees.py | 116 ++++++++++++++++++++++++ 7 files changed, 525 insertions(+) create mode 100644 docs/src/content/docs/guides/interop.md create mode 100644 tests/unit/test_interop_trees.py diff --git a/docs/astro.config.mjs b/docs/astro.config.mjs index fbf1403..3bbaa46 100644 --- a/docs/astro.config.mjs +++ b/docs/astro.config.mjs @@ -36,6 +36,7 @@ export default defineDocs({ 'guides/backends', 'guides/portable-code', 'guides/data-movement', + 'guides/interop', 'guides/array-ordering', 'guides/gpu-devices', 'guides/execution-helpers', diff --git a/docs/src/content/docs/guides/interop.md b/docs/src/content/docs/guides/interop.md new file mode 100644 index 0000000..be87a40 --- /dev/null +++ b/docs/src/content/docs/guides/interop.md @@ -0,0 +1,43 @@ +--- +title: "Interop" +description: "Hand simulation results to JAX, PyTorch or MLX and back: whole dicts, lists and tuples of arrays, and scalar objectives." +--- + +The adapters in `cunumpy.interop` (`jax`, `torch`, `mlx`) convert arrays +explicitly between NumPy or CuPy and a machine-learning framework. They share +memory where they can, never change a dtype that was not asked for, and give no +gradients through cunumpy code. Each needs its extra, e.g. `cunumpy[jax]`. + +## Nested containers + +`tree_to_framework` and `tree_from_framework` convert every array in nested +dicts, lists and tuples (named tuples included) and keep the structure; numbers, +strings and `None` pass through unchanged. The keyword arguments (`device`, +`dtype`, `copy`, `backend`) apply to each array. The JAX adapter maps over any +JAX pytree. + +```python +import numpy as np +from cunumpy.interop import jax as cj + +cj.enable_x64() +state = {"density": np.ones(64), "fields": (np.zeros(64), np.zeros(64)), "step": 10} +params = cj.tree_to_jax(state) # jax arrays, same structure, "step" stays 10 +back = cj.tree_from_jax(params) # NumPy arrays again +``` + +The aliases `tree_to_jax`, `tree_to_torch`, `tree_to_mlx` (and `tree_from_*`) +name the same functions. + +## Scalar objectives + +`scalar` returns an objective or a diagnostic as a 0-d framework array, which +is what a loss function or `jax.pure_callback` expects. It takes a Python number, +a NumPy or CuPy scalar, or any array with one element, and raises for more. +A Python float counts as float64, so the dtype rules hold: without +`jax_enable_x64` ask for the cast. + +```python +loss = cj.scalar(np.sum(back["density"] ** 2)) # 0-d float64 jax array +loss32 = cj.scalar(0.5, dtype="float32") +``` diff --git a/src/cunumpy/interop/_common.py b/src/cunumpy/interop/_common.py index 2d79afd..52e4f7d 100644 --- a/src/cunumpy/interop/_common.py +++ b/src/cunumpy/interop/_common.py @@ -432,3 +432,82 @@ def to_backend( return finish( source, result, copy=copy, target=backend, duplicate=lambda array: array.copy() ) + + +def is_array(value: Any) -> bool: + """Whether a value is an array the adapters convert (not a container). + + Parameters + ---------- + value : object + A leaf of a nested container. + + Returns + ------- + bool + True for NumPy arrays and scalars and for anything exposing DLPack or + the CUDA array interface. + """ + return isinstance(value, (np.ndarray, np.generic)) or any( + hasattr(value, name) for name in ("__dlpack__", "__cuda_array_interface__") + ) + + +def tree_map( + function: Callable[[Any], Any], tree: Any, leaf: Callable[[Any], bool] +) -> Any: + """Apply `function` to the array leaves of nested dicts, lists and tuples. + + Containers keep their type (also a named tuple and a dict subclass); a leaf + that is not an array (a Python number, a string, None) is returned as is. + + Parameters + ---------- + function : callable + Converts one array. + tree : object + An array, or dicts, lists and tuples of arrays. + leaf : callable + Tells which leaves `function` converts. + + Returns + ------- + object + The same structure with the arrays converted. + """ + if isinstance(tree, dict): + items = {key: tree_map(function, value, leaf) for key, value in tree.items()} + return items if type(tree) is dict else type(tree)(items) + if isinstance(tree, (list, tuple)): + items = [tree_map(function, value, leaf) for value in tree] + if type(tree) in (list, tuple): + return type(tree)(items) + if hasattr(tree, "_fields"): + return type(tree)(*items) + return type(tree)(items) + return function(tree) if leaf(tree) else tree + + +def scalar_value(value: Any) -> Any: + """Check that an objective or diagnostic holds one value and return it. + + Parameters + ---------- + value : object + A Python number, a NumPy or CuPy scalar, or an array of size one. + + Returns + ------- + object + `value` as a 0-d array of its library (a Python number unchanged). + """ + if isinstance(value, (bool, int, float, complex)): + return value + size = getattr(value, "size", None) + if size is None: + raise TypeError( + f"a scalar must be a number or an array, got {type(value).__name__}" + ) + if size != 1: + raise ValueError(f"a scalar must have one element, got shape {value.shape}") + return value.reshape(()) diff --git a/src/cunumpy/interop/jax.py b/src/cunumpy/interop/jax.py index f48903d..1b4eda6 100644 --- a/src/cunumpy/interop/jax.py +++ b/src/cunumpy/interop/jax.py @@ -13,8 +13,10 @@ check_dtype, device_error, finish, + is_array, memory, require, + scalar_value, to_backend, ) @@ -153,14 +155,112 @@ def from_framework( return to_backend(array, backend, dtype=dtype, copy=copy, host_array=np.asarray) +def _tree_map(function: Any, tree: Any, leaf: Any) -> Any: + """Map `function` over the leaves of a JAX pytree that `leaf` selects.""" + jax = require("jax") + return jax.tree_util.tree_map( + lambda value: function(value) if leaf(value) else value, tree + ) + + +def scalar(value: Any, *, device: Any = None, dtype: Any = None) -> Any: + """Return a scalar objective or diagnostic as a 0-d JAX array. + + A Python number becomes a NumPy scalar first (float64 for a float), so + the dtype rules of :func:`to_framework` apply to it as well. + + Parameters + ---------- + value : object + A Python number, a NumPy or CuPy scalar, or an array of size one. + device : object or None, optional + As for :func:`to_framework`. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + 0-d JAX array. + """ + value = scalar_value(value) + if not is_array(value): + value = np.asarray(value) + return to_framework(value, device=device, dtype=dtype) + + +def tree_to_framework( + tree: Any, *, device: Any = None, dtype: Any = None, copy: bool | None = None +) -> Any: + """Convert the arrays in nested dicts, lists and tuples to JAX arrays. + + The structure is kept and leaves that are not arrays (numbers, strings, + None) are returned unchanged; each array goes through + :func:`to_framework` with the same arguments. + Other JAX pytrees (registered with ``jax.tree_util``) work as well. + + Parameters + ---------- + tree : object + An array, or dicts, lists and tuples of arrays. + device, dtype, copy : object, optional + As for :func:`to_framework`. + + Returns + ------- + object + The same structure with JAX arrays. + """ + + def convert(array: Any) -> Any: + return to_framework(array, device=device, dtype=dtype, copy=copy) + + return _tree_map(convert, tree, is_array) + + +def tree_from_framework( + tree: Any, *, backend: str = "numpy", dtype: Any = None, copy: bool | None = None +) -> Any: + """Convert the JAX arrays in nested dicts, lists and tuples to NumPy or CuPy. + + The structure is kept and other leaves are returned unchanged; each JAX + array goes through :func:`from_framework` with the same arguments. + + Parameters + ---------- + tree : object + A JAX array, or dicts, lists and tuples of them. + backend, dtype, copy : object, optional + As for :func:`from_framework`. + + Returns + ------- + object + The same structure with NumPy or CuPy arrays. + """ + jax = require("jax") + + def convert(array: Any) -> Any: + return from_framework(array, backend=backend, dtype=dtype, copy=copy) + + return _tree_map(convert, tree, lambda value: isinstance(value, jax.Array)) + + to_jax = to_framework from_jax = from_framework +tree_to_jax = tree_to_framework +tree_from_jax = tree_from_framework __all__ = [ "enable_x64", "from_framework", "from_jax", + "scalar", "supports_float64", "to_framework", "to_jax", + "tree_from_framework", + "tree_from_jax", + "tree_to_framework", + "tree_to_jax", ] diff --git a/src/cunumpy/interop/mlx.py b/src/cunumpy/interop/mlx.py index 4efef0a..8c8b7d9 100644 --- a/src/cunumpy/interop/mlx.py +++ b/src/cunumpy/interop/mlx.py @@ -14,9 +14,12 @@ device_error, dtype_name, finish, + is_array, memory, require, + scalar_value, to_backend, + tree_map, ) _GPU_ADVICE = ( @@ -131,13 +134,102 @@ def from_framework( ) +def scalar(value: Any, *, device: Any = None, dtype: Any = None) -> Any: + """Return a scalar objective or diagnostic as a 0-d MLX array. + + A Python number becomes a NumPy scalar first (float64 for a float), so + the dtype rules of :func:`to_framework` apply to it as well. + + Parameters + ---------- + value : object + A Python number, a NumPy or CuPy scalar, or an array of size one. + device : object or None, optional + As for :func:`to_framework`. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + 0-d MLX array. + """ + value = scalar_value(value) + if not is_array(value): + value = np.asarray(value) + return to_framework(value, device=device, dtype=dtype) + + +def tree_to_framework( + tree: Any, *, device: Any = None, dtype: Any = None, copy: bool | None = None +) -> Any: + """Convert the arrays in nested dicts, lists and tuples to MLX arrays. + + The structure is kept and leaves that are not arrays (numbers, strings, + None) are returned unchanged; each array goes through + :func:`to_framework` with the same arguments. + + Parameters + ---------- + tree : object + An array, or dicts, lists and tuples of arrays. + device, dtype, copy : object, optional + As for :func:`to_framework`. + + Returns + ------- + object + The same structure with MLX arrays. + """ + + def convert(array: Any) -> Any: + return to_framework(array, device=device, dtype=dtype, copy=copy) + + return tree_map(convert, tree, is_array) + + +def tree_from_framework( + tree: Any, *, backend: str = "numpy", dtype: Any = None, copy: bool | None = None +) -> Any: + """Convert the MLX arrays in nested dicts, lists and tuples to NumPy or CuPy. + + The structure is kept and other leaves are returned unchanged; each MLX + array goes through :func:`from_framework` with the same arguments. + + Parameters + ---------- + tree : object + A MLX array, or dicts, lists and tuples of them. + backend, dtype, copy : object, optional + As for :func:`from_framework`. + + Returns + ------- + object + The same structure with NumPy or CuPy arrays. + """ + mx = require("mlx.core", extra="mlx") + + def convert(array: Any) -> Any: + return from_framework(array, backend=backend, dtype=dtype, copy=copy) + + return tree_map(convert, tree, lambda value: isinstance(value, mx.array)) + + to_mlx = to_framework from_mlx = from_framework +tree_to_mlx = tree_to_framework +tree_from_mlx = tree_from_framework __all__ = [ "from_framework", "from_mlx", + "scalar", "supports_float64", "to_framework", "to_mlx", + "tree_from_framework", + "tree_from_mlx", + "tree_to_framework", + "tree_to_mlx", ] diff --git a/src/cunumpy/interop/torch.py b/src/cunumpy/interop/torch.py index 51a3113..2bf645e 100644 --- a/src/cunumpy/interop/torch.py +++ b/src/cunumpy/interop/torch.py @@ -4,15 +4,20 @@ from typing import Any +import numpy as np + from cunumpy.interop._common import ( HOST, check_copy, check_dtype, device_error, finish, + is_array, memory, require, + scalar_value, to_backend, + tree_map, ) _MPS_ADVICE = ( @@ -129,13 +134,102 @@ def from_framework( ) +def scalar(value: Any, *, device: Any = None, dtype: Any = None) -> Any: + """Return a scalar objective or diagnostic as a 0-d PyTorch array. + + A Python number becomes a NumPy scalar first (float64 for a float), so + the dtype rules of :func:`to_framework` apply to it as well. + + Parameters + ---------- + value : object + A Python number, a NumPy or CuPy scalar, or an array of size one. + device : object or None, optional + As for :func:`to_framework`. + dtype : object or None, optional + Dtype for the result. + + Returns + ------- + object + 0-d PyTorch array. + """ + value = scalar_value(value) + if not is_array(value): + value = np.asarray(value) + return to_framework(value, device=device, dtype=dtype) + + +def tree_to_framework( + tree: Any, *, device: Any = None, dtype: Any = None, copy: bool | None = None +) -> Any: + """Convert the arrays in nested dicts, lists and tuples to PyTorch arrays. + + The structure is kept and leaves that are not arrays (numbers, strings, + None) are returned unchanged; each array goes through + :func:`to_framework` with the same arguments. + + Parameters + ---------- + tree : object + An array, or dicts, lists and tuples of arrays. + device, dtype, copy : object, optional + As for :func:`to_framework`. + + Returns + ------- + object + The same structure with PyTorch arrays. + """ + + def convert(array: Any) -> Any: + return to_framework(array, device=device, dtype=dtype, copy=copy) + + return tree_map(convert, tree, is_array) + + +def tree_from_framework( + tree: Any, *, backend: str = "numpy", dtype: Any = None, copy: bool | None = None +) -> Any: + """Convert the PyTorch arrays in nested dicts, lists and tuples to NumPy or CuPy. + + The structure is kept and other leaves are returned unchanged; each PyTorch + array goes through :func:`from_framework` with the same arguments. + + Parameters + ---------- + tree : object + A PyTorch array, or dicts, lists and tuples of them. + backend, dtype, copy : object, optional + As for :func:`from_framework`. + + Returns + ------- + object + The same structure with NumPy or CuPy arrays. + """ + torch = require("torch") + + def convert(array: Any) -> Any: + return from_framework(array, backend=backend, dtype=dtype, copy=copy) + + return tree_map(convert, tree, lambda value: isinstance(value, torch.Tensor)) + + to_torch = to_framework from_torch = from_framework +tree_to_torch = tree_to_framework +tree_from_torch = tree_from_framework __all__ = [ "from_framework", "from_torch", + "scalar", "supports_float64", "to_framework", "to_torch", + "tree_from_framework", + "tree_from_torch", + "tree_to_framework", + "tree_to_torch", ] diff --git a/tests/unit/test_interop_trees.py b/tests/unit/test_interop_trees.py new file mode 100644 index 0000000..9926e5e --- /dev/null +++ b/tests/unit/test_interop_trees.py @@ -0,0 +1,116 @@ +"""Tests for nested containers and scalar objectives in the framework adapters.""" + +from collections import OrderedDict, namedtuple + +import numpy as np +import pytest + +from cunumpy.interop._common import is_array, scalar_value, tree_map + +Pair = namedtuple("Pair", "a b") + + +def _tree(): + return { + "fields": [np.arange(3.0), (np.ones(2), "label")], + "pair": Pair(np.zeros(1), 2.5), + "ordered": OrderedDict(x=np.float64(1.0)), + "none": None, + } + + +def test_tree_map_keeps_structure_and_non_arrays(): + result = tree_map(lambda a: ("converted", a.shape), _tree(), is_array) + assert result["fields"][0] == ("converted", (3,)) + assert isinstance(result["fields"][1], tuple) + assert result["fields"][1][1] == "label" + assert isinstance(result["pair"], Pair) and result["pair"].b == 2.5 + assert isinstance(result["ordered"], OrderedDict) + assert result["ordered"]["x"] == ("converted", ()) + assert result["none"] is None + + +def test_tree_map_keeps_list_and_tuple_subclasses(): + class Items(list): + pass + + result = tree_map(lambda a: 0, Items([np.ones(1)]), is_array) + assert type(result) is Items and result == [0] + + +def test_scalar_value(): + assert scalar_value(1.5) == 1.5 + assert scalar_value(np.ones((1, 1))).shape == () + with pytest.raises(ValueError, match="one element"): + scalar_value(np.ones(2)) + with pytest.raises(TypeError, match="a number or an array"): + scalar_value("x") + + +@pytest.fixture +def jax_x64(): + jax = pytest.importorskip("jax") + from cunumpy.interop import jax as adapter + + enabled = jax.config.read("jax_enable_x64") + adapter.enable_x64() + yield jax, adapter + jax.config.update("jax_enable_x64", enabled) + + +def test_jax_tree_round_trip(jax_x64): + jax, adapter = jax_x64 + tree = _tree() + value = adapter.tree_to_jax(tree) + assert isinstance(value["fields"][0], jax.Array) + assert value["fields"][0].dtype == np.float64 + assert isinstance(value["pair"], Pair) and value["pair"].b == 2.5 + assert value["fields"][1][1] == "label" and value["none"] is None + back = adapter.tree_from_jax(value) + np.testing.assert_array_equal(back["fields"][0], tree["fields"][0]) + assert isinstance(back["fields"][0], np.ndarray) + assert back["pair"].b == 2.5 + + +def test_jax_tree_dtype_and_scalar(jax_x64): + jax, adapter = jax_x64 + value = adapter.tree_to_jax([np.arange(2.0)], dtype="float32") + assert value[0].dtype == np.float32 + objective = adapter.scalar(np.array([0.5])) + assert isinstance(objective, jax.Array) and objective.shape == () + assert adapter.scalar(2.0).dtype == np.float64 + assert adapter.scalar(2.0, dtype="float32").dtype == np.float32 + adapter.enable_x64(False) + with pytest.raises(TypeError, match="float64 into float32"): + adapter.scalar(2.0) + + +def test_torch_tree_round_trip_and_scalar(): + torch = pytest.importorskip("torch") + from cunumpy.interop import torch as adapter + + tree = _tree() + value = adapter.tree_to_torch(tree) + assert isinstance(value["fields"][0], torch.Tensor) + assert value["fields"][0].dtype == torch.float64 + assert isinstance(value["pair"], Pair) and value["fields"][1][1] == "label" + tensor = torch.ones(2, requires_grad=True) + back = adapter.tree_from_torch({"t": tensor, "n": 3}) + assert isinstance(back["t"], np.ndarray) and back["n"] == 3 + objective = adapter.scalar(np.float64(1.5)) + assert objective.shape == () and objective.dtype == torch.float64 + assert adapter.scalar(1, dtype=torch.float32).dtype == torch.float32 + + +def test_mlx_tree_round_trip_and_scalar(): + mx = pytest.importorskip("mlx.core") + from cunumpy.interop import mlx as adapter + + value = adapter.tree_to_mlx( + {"a": np.arange(3, dtype=np.float32), "s": "x"}, device=mx.cpu + ) + assert isinstance(value["a"], mx.array) and value["s"] == "x" + back = adapter.tree_from_mlx(value) + np.testing.assert_array_equal(back["a"], np.arange(3, dtype=np.float32)) + objective = adapter.scalar(1.5, device=mx.cpu) + assert objective.shape == () and objective.dtype == mx.float64