From d6d1f70850ab69c61936cf206199b2617edcbe4a Mon Sep 17 00:00:00 2001 From: crusaderky Date: Sun, 26 Jan 2025 19:12:50 +0000 Subject: [PATCH] ENH: ndonnx device() support; TST: better ndonnx test coverage --- array_api_compat/common/_helpers.py | 4 ++-- docs/supported-array-libraries.md | 5 +++++ tests/_helpers.py | 7 ++++--- tests/test_array_namespace.py | 3 +++ tests/test_common.py | 12 +++++++++++- 5 files changed, 25 insertions(+), 6 deletions(-) diff --git a/array_api_compat/common/_helpers.py b/array_api_compat/common/_helpers.py index 434e7d87..966b6e1a 100644 --- a/array_api_compat/common/_helpers.py +++ b/array_api_compat/common/_helpers.py @@ -645,7 +645,7 @@ def device(x: Array, /) -> Device: to_device : Move array data to a different device. """ - if is_numpy_array(x): + if is_numpy_array(x) or is_ndonnx_array(x): return "cpu" elif is_dask_array(x): # Peek at the metadata of the Dask array to determine type @@ -776,7 +776,7 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]] device : Hardware device the array data resides on. """ - if is_numpy_array(x): + if is_numpy_array(x) or is_ndonnx_array(x): if stream is not None: raise ValueError("The stream argument to to_device() is not supported") if device == 'cpu': diff --git a/docs/supported-array-libraries.md b/docs/supported-array-libraries.md index 1af9f3dc..4519c4ac 100644 --- a/docs/supported-array-libraries.md +++ b/docs/supported-array-libraries.md @@ -138,6 +138,11 @@ The minimum supported Dask version is 2023.12.0. Similar to JAX, `sparse` Array API support is contained directly in `sparse`. +(ndonnx-support)= +## [ndonnx](https://github.com/quantco/ndonnx) + +Similar to JAX, `ndonnx` Array API support is contained directly in `ndonnx`. + (array-api-strict-support)= ## [array-api-strict](https://data-apis.org/array-api-strict/) diff --git a/tests/_helpers.py b/tests/_helpers.py index 5b79aa46..82623f87 100644 --- a/tests/_helpers.py +++ b/tests/_helpers.py @@ -3,11 +3,12 @@ import pytest wrapped_libraries = ["numpy", "cupy", "torch", "dask.array"] -all_libraries = wrapped_libraries + ["array_api_strict", "jax.numpy", "sparse"] - +all_libraries = wrapped_libraries + [ + "array_api_strict", "jax.numpy", "ndonnx", "sparse" +] def import_(library, wrapper=False): - if library == 'cupy': + if library in ('cupy', 'ndonnx'): pytest.importorskip(library) if wrapper: if 'jax' in library: diff --git a/tests/test_array_namespace.py b/tests/test_array_namespace.py index a66a64d9..605c69a1 100644 --- a/tests/test_array_namespace.py +++ b/tests/test_array_namespace.py @@ -22,6 +22,9 @@ def test_array_namespace(library, api_version, use_compat): if use_compat and library not in wrapped_libraries: pytest.raises(ValueError, lambda: array_namespace(array, use_compat=use_compat)) return + if library == "ndonnx" and api_version in ("2021.12", "2022.12"): + pytest.skip("Unsupported API version") + namespace = array_namespace(array, api_version=api_version, use_compat=use_compat) if use_compat is False or use_compat is None and library not in wrapped_libraries: diff --git a/tests/test_common.py b/tests/test_common.py index e95e305e..5bc8e98a 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -8,9 +8,10 @@ from array_api_compat import ( # noqa: F401 is_numpy_array, is_cupy_array, is_torch_array, is_dask_array, is_jax_array, is_pydata_sparse_array, + is_ndonnx_array, is_numpy_namespace, is_cupy_namespace, is_torch_namespace, is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace, - is_array_api_strict_namespace, + is_array_api_strict_namespace, is_ndonnx_namespace, ) from array_api_compat import ( @@ -25,6 +26,7 @@ 'dask.array': 'is_dask_array', 'jax.numpy': 'is_jax_array', 'sparse': 'is_pydata_sparse_array', + 'ndonnx': 'is_ndonnx_array', } is_namespace_functions = { @@ -35,6 +37,7 @@ 'jax.numpy': 'is_jax_namespace', 'sparse': 'is_pydata_sparse_namespace', 'array_api_strict': 'is_array_api_strict_namespace', + 'ndonnx': 'is_ndonnx_namespace', } @@ -232,6 +235,13 @@ def _xfail(reason: str) -> None: # TODO: remove xfail once # https://github.com/dask/dask/issues/8260 is resolved _xfail(reason="Bug in dask raising error on conversion") + elif ( + source_library == "ndonnx" + and target_library not in ("array_api_strict", "ndonnx", "numpy") + ): + _xfail(reason="The truth value of lazy Array Array(dtype=Boolean) is unknown") + elif source_library == "ndonnx" and target_library == "numpy": + _xfail(reason="produces numpy array of ndonnx scalar arrays") elif source_library == "jax.numpy" and target_library == "torch": _xfail(reason="casts int to float") elif source_library == "cupy" and target_library != "cupy":