ndonnx support

Author: crusaderky

array_api_compat/common/_helpers.py

@@ -627,17 +627,13 @@ 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
-        try:
-            import numpy as np
-            if isinstance(x._meta, np.ndarray):
-                # Must be on CPU since backed by numpy
-                return "cpu"
-        except ImportError:
-            pass
+        if is_numpy_array(x._meta):
+            # Must be on CPU since backed by numpy
+            return "cpu"
         return _DASK_DEVICE
     elif is_jax_array(x):
         # JAX has .device() as a method, but it is being deprecated so that it
@@ -758,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':
@@ -780,7 +776,6 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
         if not hasattr(x, "__array_namespace__"):
             # In JAX v0.4.31 and older, this import adds to_device method to x.
             import jax.experimental.array_api # noqa: F401
-        return x.to_device(device, stream=stream)
     elif is_pydata_sparse_array(x) and device == _device(x):
         # Perform trivial check to return the same array if
         # device is same instead of err-ing.

tests/test_array_namespace.py

@@ -14,14 +14,17 @@
 
 @pytest.mark.parametrize("use_compat", [True, False, None])
 @pytest.mark.parametrize("api_version", [None, "2021.12", "2022.12", "2023.12"])
-@pytest.mark.parametrize("library", all_libraries + ['array_api_strict'])
+@pytest.mark.parametrize("library", all_libraries)
 def test_array_namespace(library, api_version, use_compat):
     xp = import_(library)
 
     array = xp.asarray([1.0, 2.0, 3.0])
-    if use_compat is True and library in {'array_api_strict', 'jax.numpy', 'sparse'}:
+    if use_compat and library in {'array_api_strict', 'jax.numpy', 'ndonnx', 'sparse'}:
         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

@@ -144,12 +144,21 @@ def test_asarray_cross_library(source_library, target_library, request):
 
         # TODO: remove xfail once
         # https://github.com/dask/dask/issues/8260 is resolved
-        request.node.add_marker(pytest.mark.xfail(reason="Bug in dask raising error on conversion"))
-    if source_library == "cupy" and target_library != "cupy":
+        request.node.add_marker(
+            pytest.mark.xfail(reason="Bug in dask raising error on conversion")
+        )
+    elif source_library == "ndonnx" and target_library not in ("numpy", "array_api_strict"):
+        request.node.add_marker(
+            pytest.mark.xfail(
+                reason="The truth value of lazy Array Array(dtype=Boolean) is unknown"
+            )
+        )        
+    elif source_library == "cupy" and target_library != "cupy":
         # cupy explicitly disallows implicit conversions to CPU
         pytest.skip(reason="cupy does not support implicit conversion to CPU")
     elif source_library == "sparse" and target_library != "sparse":
         pytest.skip(reason="`sparse` does not allow implicit densification")
+
     src_lib = import_(source_library, wrapper=True)
     tgt_lib = import_(target_library, wrapper=True)
     is_tgt_type = globals()[is_array_functions[target_library]]