Interop: nested containers and scalar objectives

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

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")
+```

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(())

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

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