ENH: ndonnx device() support; TST: better ndonnx test coverage

Author: crusaderky

array_api_compat/common/_helpers.py

@@ -627,7 +627,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 jax array to determine type
@@ -754,7 +754,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':

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

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:

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_api_compat.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:

tests/test_common.py

@@ -1,9 +1,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 (
@@ -23,6 +24,7 @@
     'dask.array': 'is_dask_array',
     'jax.numpy': 'is_jax_array',
     'sparse': 'is_pydata_sparse_array',
+    'ndonnx': 'is_ndonnx_array',
 }
 
 is_namespace_functions = {
@@ -33,6 +35,7 @@
     'jax.numpy': 'is_jax_namespace',
     'sparse': 'is_pydata_sparse_namespace',
     'array_api_strict': 'is_array_api_strict_namespace',
+    'ndonnx': 'is_ndonnx_namespace',
 }
 
 
@@ -167,6 +170,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":