ENH: ndonnx device() support; TST: better ndonnx test coverage #232
Conversation
crusaderky
commented
Jan 8, 2025
•
crusaderky
commented
- Add support to ndonnx in device() and to_device().
- Add ndonnx to parameterized tests
- Add ndonnx to the documentation
- Adds key use cases to ENH: is_lazy_array() #228
- Adds key use cases to ENH: size() to return None on dask instead of nan #231
|
This is now ready for review |
crusaderky
force-pushed
the
ndonnx
branch
3 times, most recently
from
48ed46c to
8cac708
Compare
crusaderky
force-pushed
the
ndonnx
branch
2 times, most recently
from
08d5270 to
fc4079f
Compare
There was a problem hiding this comment.
What concerns me here is a sort of conceptual question: What is the scope of a test run here in array-api-compat?
Basically, there are two kinds of array libraries: those which are array API compatible and those which need some wrapping. array-api-compat does the latter, for those libraries which need it.
Than why are we testing the libraries which we do not wrap?
There's no question that we need some better ergonomics for running tests across array libraries. I'm just not convinced it's a good idea to sort of sneak this testing here. Let's discuss a structural solution instead.
| @@ -772,7 +772,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): | |||
There was a problem hiding this comment.
What I don't quite understand is this: https://github.com/data-apis/array-api-compat/pull/232/files#diff-d5d8fa69860e03769cc235a54027efad4376a8c48eeab71aaa11365ea308123bR141 states that the array API support is internal to ndonnx, which seems to imply there's no need to special-case it here?
The note reads "Similar to JAX, ndonnx Array API support is contained directly in ndonnx." --- and indeed, there are no workarounds for jax in the codebase.
There was a problem hiding this comment.
Because they support the standard imperfectly. The array-api-tests for .device and .to_device are XFAILed:
https://github.com/Quantco/ndonnx/actions/workflows/array-api.yml
api-coverage-tests/array_api_tests/test_has_names.py::test_has_names[array_method-to_device]
[gw1] [ 13%] XFAIL api-coverage-tests/array_api_tests/test_has_names.py::test_has_names[array_method-to_device]
api-coverage-tests/array_api_tests/test_has_names.py::test_has_names[array_attribute-T]
[gw1] [ 14%] PASSED api-coverage-tests/array_api_tests/test_has_names.py::test_has_names[array_attribute-T]
api-coverage-tests/array_api_tests/test_has_names.py::test_has_names[array_attribute-device]
[gw1] [ 14%] XFAIL api-coverage-tests/array_api_tests/test_has_names.py::test_has_names[array_attribute-device]
Adding .device and .to_device won't break any library's backwards compatibility, so it would be best if each library just fixed it on their side. And yet here we have a function that works around the quirks of every library - last but not least to support least-than-newest versions of it.
There was a problem hiding this comment.
I will address this upstream. I should say that it might be deceptive to allow this to work with "cpu".
A serialized ONNX graph from ndonnx could later be executed on many different hardware targets or runtimes. ONNX is one level removed from the execution target by design.
Based on a quick read of the device page of the specification, we may do something similar to what you're doing with Dask.
| all_libraries = wrapped_libraries + ["array_api_strict", "jax.numpy", "sparse"] | ||
|
|
||
| all_libraries = wrapped_libraries + [ | ||
| "array_api_strict", "jax.numpy", "ndonnx", "sparse" |
There was a problem hiding this comment.
This too. Why are we testing here libraries which we don't wrap?
There was a problem hiding this comment.
Because _helpers.py is chock full of special cases and exceptions for them.
because the libraries that we do not wrap still need the functions defined in |
There are 66 references to JAX in |
Yes, most of which are |
|
@crusaderky we added device support in Quantco/ndonnx@039264d. It should be possible to remove the device support here once we make a release. Is it urgent that we land this PR soon? |
No, this PR can wait. |
