Replace `int` with `SupportsIndex` in indexing methods hints

data-apis / array-api

RFC document, tooling and other content related to the array API standard

https://data-apis.github.io/array-api/latest/

MIT License

204 stars 42 forks source link

Replace `int` with `SupportsIndex` in indexing methods hints #766

Open honno opened 3 months ago

honno commented 3 months ago

Resolves https://github.com/data-apis/array-api/issues/383

rgommers commented 3 months ago

Thanks @honno.

Adding a quick clarifying example:

>>> import numpy as np
>>> import torch
>>> 
>>> x = np.arange(1, 4)
>>> y = torch.arange(1, 4)
>>> 
>>> class Ix:
...     def __index__(self):
...         return 1
... 
>>> x[Ix()]
2
>>> y[Ix()]
tensor(2)

rgommers commented 3 months ago

@honno would you be able to open a PR with a test for this to array-api-tests? Doesn't have to be merged before this will be merged, but that will help smoke out if any known/tested library does not yet support this feature.

honno commented 3 months ago

@honno would you be able to open a PR with a test for this to array-api-tests? Doesn't have to be merged before this will be merged, but that will help smoke out if any known/tested library does not yet support this feature.

Good shout, I opened https://github.com/data-apis/array-api-tests/pull/247 to check this all out. From the looks of it:

NumPy (+ array_api_strict) and PyTorch supports this behaviour.
JAX doesn't support, has a general note to open a feature requests on indexing modes (... that's prob refering to advance integer indexing).
cupy seems to only support advance integer indexing and doesn't play nice with 0-D integer arrays as indexes.
Dask at least accepts 0-D integer input, but array_api_compat.dask not playing nice with the test suite so I'll have to triage that and explore further.

asmeurer commented 3 months ago

but array_api_compat.dask not playing nice with the test suite so I'll have to triage that and explore further.

Let me know what you find out. We do run it on CI with some skips and xfails, and also the max-examples is set to 5. I haven't looked at the Dask xfails too closely, and obviously if we can remove any of those that would be great. CC @lithomas1

kgryte commented 2 months ago

@honno Were you able to triage the Dask issues with the test suite?

honno commented 2 months ago

Got Dask working[^1]—the latest release at least doesn't support indexables for both get and set items right now (TypeError for an internal inequality check which assumes indexes as ints).

[^1]: Turns out I should of been using array_api_compat.dask.array heh, although there's another unrelated issue Dask has with one of our utilities I'll have to explore.

asmeurer commented 2 months ago

By the way, for implementers, the generally correct behavior is to operator.index() to normalize index objects that aren't one of the other supported index types like Ellipsis, slice, or array. My guess is that dask.array isn't doing that.