8000 ENH: ndonnx device() support; TST: better ndonnx test coverage · data-apis/array-api-compat@48ed46c · GitHub
[go: up one dir, main page]
Skip to content

Commit 48ed46c

Browse files
crusaderkycrusaderky
committed
ENH: ndonnx device() support; TST: better ndonnx test coverage
1 parent 624ae63 commit 48ed46c

File tree

5 files changed

+25
-6
lines changed

5 files changed

+25
-6
lines changed

‎array_api_compat/common/_helpers.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -627,7 +627,7 @@ def device(x: Array, /) -> Device:
627627
to_device : Move array data to a different device.
628628
629629
"""
630-
if is_numpy_array(x):
630+
if is_numpy_array(x) or is_ndonnx_array(x):
631631
return "cpu"
632632
elif is_dask_array(x):
633633
# 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]]
754754
device : Hardware device the array data resides on.
755755
756756
"""
757-
if is_numpy_array(x):
757+
if is_numpy_array(x) or is_ndonnx_array(x):
758758
if stream is not None:
759759
raise ValueError("The stream argument to to_device() is not supported")
760760
if device == 'cpu':

‎docs/supported-array-libraries.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -138,6 +138,11 @@ The minimum supported Dask version is 2023.12.0.
138138

139139
Similar to JAX, `sparse` Array API support is contained directly in `sparse`.
140140

141+
(ndonnx-support)=
142+
## [ndonnx](https://github.com/quantco/ndonnx)
143+
144+
Similar to JAX, `ndonnx` Array API support is contained directly in `ndonnx`.
145+
141146
(array-api-strict-support)=
142147
## [array-api-strict](https://data-apis.org/array-api-strict/)
143148

‎tests/_helpers.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,12 @@
33
import pytest
44

55
wrapped_libraries = ["numpy", "cupy", "torch", "dask.array"]
6-
all_libraries = wrapped_libraries + ["array_api_strict", "jax.numpy", "sparse"]
7-
6+
all_libraries = wrapped_libraries + [
7+
"array_api_strict", "jax.numpy", "ndonnx", "sparse"
8+
]
89

910
def import_(library, wrapper=False):
10-
if library == 'cupy':
11+
if library in ('cupy', 'ndonnx'):
1112
pytest.importorskip(library)
1213
if wrapper:
1314
if 'jax' in library:

‎tests/test_array_namespace.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,9 @@ def test_array_namespace(library, api_version, use_compat):
2222
if use_compat and library not in wrapped_libraries:
2323
pytest.raises(ValueError, lambda: array_namespace(array, use_compat=use_compat))
2424
return
25+
if library == "ndonnx" and api_version in ("2021.12", "2022.12"):
26+
pytest.skip("Unsupported API version")
27+
2528
namespace = array_api_compat.array_namespace(array, api_version=api_version, use_compat=use_compat)
2629

2730
if use_compat is False or use_compat is None and library not in wrapped_libraries:

‎tests/test_common.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,10 @@
88
from array_api_compat import ( # noqa: F401
99
is_numpy_array, is_cupy_array, is_torch_array,
1010
is_dask_array, is_jax_array, is_pydata_sparse_array,
11+
is_ndonnx_array,
1112
is_numpy_namespace, is_cupy_namespace, is_torch_namespace,
1213
is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace,
13-
is_array_api_strict_namespace,
14+
is_array_api_strict_namespace, is_ndonnx_namespace,
1415
)
1516

1617
from array_api_compat import (
@@ -25,6 +26,7 @@
2526
'dask.array': 'is_dask_array',
2627
'jax.numpy': 'is_jax_array',
2728
'sparse': 'is_pydata_sparse_array',
29+
'ndonnx': 'is_ndonnx_array',
2830
}
2931

3032
is_namespace_functions = {
@@ -35,6 +37,7 @@
3537
'jax.numpy': 'is_jax_namespace',
3638
'sparse': 'is_pydata_sparse_namespace',
3739
'array_api_strict': 'is_array_api_strict_namespace',
40+
'ndonnx': 'is_ndonnx_namespace',
3841
}
3942

4043

@@ -208,6 +211,13 @@ def _xfail(reason: str) -> None:
208211
# TODO: remove xfail once
209212
# https://github.com/dask/dask/issues/8260 is resolved
210213
_xfail(reason="Bug in dask raising error on conversion")
214+
elif (
215+
source_library == "ndonnx"
216+
and target_library not in ("array_api_strict", "ndonnx", "numpy")
217+
):
218+
_xfail(reason="The truth value of lazy Array Array(dtype=Boolean) is unknown")
219+
elif source_library == "ndonnx" and target_library == "numpy":
220+
_xfail(reason="produces numpy array of ndonnx scalar arrays")
211221
elif source_library == "jax.numpy" and target_library == "torch":
212222
_xfail(reason="casts int to float")
213223
elif source_library == "cupy" and target_library != "cupy":

0 commit comments

Comments
 (0)
0