Python Array API Standard

[TomAugspurger]

Hi all,

A group of folks from various projects / companies have put together a draft of an Array API standard, and are looking for feedback.

• Blog post: https://data-apis.org/blog/array_api_standard_release
• Array API standard document: https://data-apis.github.io/array-api/latest/
• Repo: https://github.com/data-apis/array-api/

We have lots of experience working with Arrays, both as consumers of arrays implementing the NumPy API (cupy, sparse) and implementing the API ourselves. Our feedback on the proposed standard would be welcome in the GitHub repo linked above.

We can use this issue to discuss if and how dask.array should adopt this standard. The document gives some guidance at https://data-apis.github.io/array-api/latest/purpose_and_scope.html#how-to-adopt-this-api, but this is fairly uncertain territory.

[jakirkham]

Do we have someone representing us on the array side? If not, I'd be happy to volunteer. Though would need to get plugged into the meetings somehow

[kkraus14]

cc @rgommers for visibility

[rgommers]

Thanks @jakirkham! So far it's been mainly @TomAugspurger and @shoyer who contributed Dask expertise. Would be great to have you join. Also happy to have a call soon to catch you up on what's been happening. I'll try and collect a few of the Dask-related points (there weren't many at all in the past few months) and will get back to you.

Also, let me point out the NumPy Enhancement Proposal for adoption of the array API standard: numpy/numpy#18456

[jakirkham]

Sounds good. Happy to do some reading as well (guessing there's lots to catch up on). That would be great. Thank you Ralf! 😄

Yep saw that. Will take a look at that as well 🙂

[rgommers]

The one issue with an unresolved Dask-related angle is data-apis/array-api#97 (related to data/value-dependent output shapes and how to represent "unknown shape" - Dask uses nan and TensorFlow None).

data-apis/array-api#84 on including boolean indexing that's hard for Dask to support may also be good to review. Discussed concluded with color-coding that indexing mode as "not all libraries may support this, it's hard in combination with delayed evaluation and JIT compilation" (not implemented yet, that's why the issue is still open).

Other than that, I'd suggest looking at https://data-apis.github.io/array-api/latest/ and https://numpy.org/neps/nep-0047-array-api-standard.html for things that are suboptimal or missing from Dask's perspective.

[jakirkham]

Commented on the nan issue. For context this was added in PR ( dask/dask#1838 ). Maybe it can change, but to set expectations it is likely a decent lift and error prone change. Additionally other maintainers may be wary of such a change as it could be a source of bugs going forward. Admittedly maybe other libraries feel similarly depending on how heavily they use this feature. Dask is fairly reliant on it as you have seen with masking, raveling, or other operations where the content of the array may affect the result (unique).

Dask does support a fair bit of boolean indexing (though we may find some edge cases that don't work). However as the mask itself may have unknown values, this falls in the category of things where the array's content affects the array shape and so the shape of the array is unknown. I'm not sure I'd write this off as not supported, but it's also possible I don't understand what supported would mean in the context of the spec or the cases you have in mind. So more context here would be welcome

Ok will read up on this before we talk more 🙂

---

Edit: It would be good to discuss ravel somewhere. Have held off opening an issue, but I think trying to match NumPy too closely here is a pain point for Dask. If we could come up with a sensible way to relax that, which we agree would be acceptable, some (though certainly not all) unknown shape issues would go away. This would also offer a performance boost for users. See issue ( dask/dask#4855 ) for more details

[jakirkham]

Also related to Data API framework. Generally we reach for and write functions to check if they are generally an Array, DataFrame, Series, etc. ( dask/dask#7327 ) Not sure if there are general tools scoped out in this framework to make those checks easier, but that would be very helpful for Dask

[leofang]

Interesting. AFAIK there is not yet a way to distinguish a Data API array from a Data API dataframe. But since the latter is still in a draft status (more "drafty" than the former, as it's not yet made public), perhaps we can discuss about it?

[rgommers]

Edit: It would be good to discuss ravel somewhere.

There is no ravel function. The NumPy API is messy in this regard: ravel, flatten, flat, and reshape are all similar, with some random differences like returning a copy vs. a view (which cannot be a distinguishing feature for the API standard, views are an implementation detail). The standard just has reshape; reshape(-1) is equivalent to ravel.

Generally we reach for and write functions to check if they are generally an Array, DataFrame, Series, etc. ( dask/dask#7327 ) Not sure if there are general tools scoped out in this framework to make those checks easier, but that would be very helpful for Dask

Arrays which support the standard, or from which you can at least retrieve the standard namespace with a compliant implementation, have a __array_namespace__ attribute. That should be the easiest way to detect an array. For static typing we may develop a protocol (see https://data-apis.org/array-api/latest/design_topics/static_typing.html), but given that there's no plan to have a shared reference library that other libraries can depend on, there's no good way to design an isinstance check. Which is fine I think, a simple utility function that uses hasattr(x, '__array_namespace__') should be fine.

[leofang]

OK. With the ongoing discussion with the dataframe API, a similar check can be done for dataframes: hasattr(x, '__dataframe__'). But, it's important to note that neither attributes __array_namespace__ and __dataframe__ can have any side effect, or it'd be triggered when calling hasattr. Not sure if it's possible to enforce this requirement in the Array/Dataframe API Standard...

[rgommers]

But, it's important to note that neither attributes __array_namespace__ and __dataframe__ can have any side effect, or it'd be triggered when calling hasattr. Not sure if it's possible to enforce this requirement in the Array/Dataframe API Standard...

Already enforced by them being methods I think - attributes can easily have side effects, methods can't. Of course Python is dynamic in nature, so you could come up with some hack that has an unwanted side effect, but it's highly unlikely to occur in practice.

[honno]

If folks wanted to start using the compliance suite array-api-tests, you might like to check out this simple wrapper that implements some crucial things most tests require. I wasn't sure if dask.array was also going to be your Array API namespace ala PyTorch—that's what I assumed with this at least. https://gist.github.com/honno/70623811b88c0f5b224c668ec4ed86b7

Generally, please feel free to reach out via issues for any questions/bugs. We want to make the suite easier to use for implementors, so any feedback would be appreciated.

[tomwhite]

Thanks @honno for providing the wrapper and the test suite.

There are likely to be some compatibility issues with using dask.array as the Array API namespace. To pick one example, in Dask the signature for eye is

eye(N, chunks='auto', M=None, k=0, dtype=<class 'float'>)

whereas in the Array API it is

eye(n_rows, n_cols, /, *, k=0  dtype=None, device=None)

This means that an expression like eye(100, 10) means two different things: in Dask 10 refers to the chunk size, but in the Array API it's the number of columns (M). It's not clear if these can be made compatible.

It might be sensible therefore to use a different module, like dask.array_api, just like NumPy, CuPy, and MXNet do. As a proof of concept I've implemented a wrapper to do that in this branch. (File naming and structure is based on numpy.array_api.) The array object is a regular dask.array.Array though.

Running the test suite I get:

ARRAY_API_TESTS_MODULE=dask.array_api pytest --max-examples=2
[snip]
====================================================================== short test summary info =======================================================================
FAILED array_api_tests/test_array_object.py::test_getitem_masking - IndexError: Too many indices for array
FAILED array_api_tests/test_creation_functions.py::test_eye - TypeError: dtype must be known for auto-chunking
FAILED array_api_tests/test_creation_functions.py::test_linspace - ZeroDivisionError: float division by zero
FAILED array_api_tests/test_data_type_functions.py::test_can_cast - AssertionError: out=True, but should be False [can_cast(bool, int8)]
FAILED array_api_tests/test_linalg.py::test_tensordot - AssertionError: out.dtype=uint64, but should be uint8 [tensordot(uint8, uint8)]
FAILED array_api_tests/test_manipulation_functions.py::test_concat - TypeError: '<' not supported between instances of 'NoneType' and 'int'
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_positive[positive] - IndexError: Too many indices for array
FAILED array_api_tests/test_operators_and_elementwise_functions.py::test_positive[__pos__] - IndexError: Too many indices for array
FAILED array_api_tests/test_searching_functions.py::test_argmax - ValueError: At least one iterator operand must be non-NULL
FAILED array_api_tests/test_searching_functions.py::test_argmin - ValueError: At least one iterator operand must be non-NULL
FAILED array_api_tests/test_searching_functions.py::test_nonzero - AssertionError: out[0].size=1, but should be out[0].size=nan
FAILED array_api_tests/test_set_functions.py::test_unique_all - AssertionError: out.indices.shape=(nan,), but should be out.values.shape=(nan,)
FAILED array_api_tests/test_set_functions.py::test_unique_counts - AssertionError: out.counts.shape=(nan,), but should be out.values.shape=(nan,)
FAILED array_api_tests/test_set_functions.py::test_unique_inverse - TypeError: 'float' object cannot be interpreted as an integer
FAILED array_api_tests/test_set_functions.py::test_unique_values - TypeError: 'float' object cannot be interpreted as an integer
FAILED array_api_tests/test_signatures.py::test_has_names[__array_namespace__] - AssertionError: The array object is missing the method __array_namespace__()
FAILED array_api_tests/test_signatures.py::test_function_positional_args[eye] - AssertionError: Unexpected exception TypeError('dtype must be known for auto-chunki...
FAILED array_api_tests/test_signatures.py::test_function_positional_args[matmul] - TypeError: dtype must be known for auto-chunking
FAILED array_api_tests/test_signatures.py::test_function_positional_args[matrix_transpose] - TypeError: dtype must be known for auto-chunking
FAILED array_api_tests/test_signatures.py::test_function_positional_args[tensordot] - TypeError: dtype must be known for auto-chunking
FAILED array_api_tests/test_signatures.py::test_function_positional_args[vecdot] - TypeError: dtype must be known for auto-chunking
FAILED array_api_tests/test_signatures.py::test_function_keyword_only_args[eye] - AssertionError: Unexpected exception TypeError('dtype must be known for auto-chun...
FAILED array_api_tests/test_signatures.py::test_function_keyword_only_args[matmul] - TypeError: dtype must be known for auto-chunking
FAILED array_api_tests/test_signatures.py::test_function_keyword_only_args[matrix_transpose] - TypeError: dtype must be known for auto-chunking
FAILED array_api_tests/test_signatures.py::test_function_keyword_only_args[tensordot] - TypeError: dtype must be known for auto-chunking
FAILED array_api_tests/test_signatures.py::test_function_keyword_only_args[vecdot] - TypeError: dtype must be known for auto-chunking
================================================ 26 failed, 3544 passed, 541 skipped, 5 warnings in 83.74s (0:01:23) =================================================

So: a few failures, as might be expected, but also lots of tests passing, which is very encouraging.

A few caveats: I ran with only a maximum of two hypothesis examples just to do a quick check at this stage. Increasing it would almost certainly find more failing cases. I also used the following skips.txt to exclude functions that are known to be missing in Dask (and also special cases which are failing).

# expand_dims is not implemented in Dask
array_api_tests/test_manipulation_functions.py::test_expand_dims
array_api_tests/test_signatures.py::test_has_names[expand_dims]

# sorting is not implemented in Dask (https://github.com/dask/dask/issues/4368)
array_api_tests/test_sorting_functions.py::test_argsort
array_api_tests/test_sorting_functions.py::test_sort
array_api_tests/test_signatures.py::test_has_names[argsort]
array_api_tests/test_signatures.py::test_has_names[sort]

# in-place array operations are not implemented in Dask (https://github.com/dask/dask/issues/5199)
array_api_tests/test_signatures.py::test_has_names[__iadd__]
array_api_tests/test_signatures.py::test_has_names[__iand__]
array_api_tests/test_signatures.py::test_has_names[__ifloordiv__]
array_api_tests/test_signatures.py::test_has_names[__ilshift__]
array_api_tests/test_signatures.py::test_has_names[__imatmul__]
array_api_tests/test_signatures.py::test_has_names[__imod__]
array_api_tests/test_signatures.py::test_has_names[__imul__]
array_api_tests/test_signatures.py::test_has_names[__ior__]
array_api_tests/test_signatures.py::test_has_names[__ipow__]
array_api_tests/test_signatures.py::test_has_names[__irshift__]
array_api_tests/test_signatures.py::test_has_names[__isub__]
array_api_tests/test_signatures.py::test_has_names[__itruediv__]
array_api_tests/test_signatures.py::test_has_names[__ixor__]

# DLPack is not implemented in Dask
array_api_tests/test_signatures.py::test_has_names[__dlpack__]
array_api_tests/test_signatures.py::test_has_names[__dlpack_device__]
array_api_tests/test_signatures.py::test_has_names[from_dlpack]

# device support is not implemented in Dask
array_api_tests/test_signatures.py::test_has_names[to_device]
array_api_tests/test_signatures.py::test_has_names[device]

# don't run special cases yet
array_api_tests/special_cases

The next steps might be to investigate the failures, and to implement some of the functions that are missing in Dask (some definitely easier than others).

[jakirkham]

Thanks for pushing on this @tomwhite! 😄

Nice to see some initial success here. Agree we would want a different array_api module like other libraries have. Just like the other libraries we wouldn't want to break existing functionality and would want users to opt-in at some level.

Yeah this is a good point that really hasn't come up in our discussions thus far @rgommers @kgryte (likely because it is very Dask specific), but creation functions in Dask need a chunks argument. There really isn't something like this in the Array API and I'm not sure it makes much sense for other libraries, but perhaps we could allow these functions to take some additional arguments that the spec leaves open ended? Anyways would be good to get your thoughts on this.

Looking at the failures. Guessing getitem exercises some selection case that Dask doesn't implement (though would be good to know what it is and link that to a Dask issue). With eye, probably the Array API needs to skip passing dtype for arguments that are None or Dask should replace None with the default type internally. Same with other functions that see similar issues. linspace smells like a bug (would just need to look into the input arguments and file a related issue). Dask doesn't implement can_cast so it appears that failure comes from NumPy. Dask's tensordot promotes types. unique has nan in shape for unknown dimensions (don't really see that changing). Though the Array API wrapper could translate nan to None (would require some kind of wrapper Array object most likely). Not really sure with the other error messages (probably require deeper dives).

Seems we fallback to auto chunking in a lot of cases. Idk how desirable that is. It is simpler to implement, but may have performance costs depending on what chunks are needed and what are used.