Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support NumPy array API (experimental) #6804

Merged
merged 12 commits into from
Jul 20, 2022
Merged

Conversation

tomwhite
Copy link
Contributor

@tomwhite tomwhite commented Jul 18, 2022

Note: I haven't actually tested this with pytorch (which is the motivating example for #3232).

if isinstance(data, NON_NUMPY_SUPPORTED_ARRAY_TYPES):
if (
isinstance(data, NON_NUMPY_SUPPORTED_ARRAY_TYPES)
or hasattr(data, "__array_function__")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's an __array_function__ check down below (lines 245, so it seems like this change should go there).

I suspect we can also delete cupy, dask from NON_NUMPY_SUPPORTED_ARRAY_TYPES

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There's an __array_function__ check down below (lines 245, so it seems like this change should go there).

Good point - thanks for the suggestion! I've moved the change to down below.

I suspect we can also delete cupy, dask from NON_NUMPY_SUPPORTED_ARRAY_TYPES

I haven't tried this. Sounds like this would be a separate change?

value = value[(slice(None),) * axis + (subkey, Ellipsis)]
return value
else:
assert isinstance(key, VectorizedIndexer)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit pick: since assert statements are removed when python is invoked with -O and -OO parameters, could we raise a proper exception?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes we should raise I think

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Replaced with a TypeError

Comment on lines 1321 to 1324
if isinstance(key, BasicIndexer):
self.array[key.tuple] = value
elif isinstance(key, OuterIndexer):
self.array[key.tuple] = value
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

unless i'm missing something, we should be able to combine these two branches into one, correct?

Suggested change
if isinstance(key, BasicIndexer):
self.array[key.tuple] = value
elif isinstance(key, OuterIndexer):
self.array[key.tuple] = value
if isinstance(key, (BasicIndexer, OuterIndexer)):
self.array[key.tuple] = value

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct - fixed.

@TomNicholas TomNicholas added the topic-arrays related to flexible array support label Jul 18, 2022
xarray/core/indexing.py Outdated Show resolved Hide resolved
xarray/core/indexing.py Outdated Show resolved Hide resolved
tomwhite and others added 2 commits July 19, 2022 12:03
Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
Copy link
Contributor

@Illviljan Illviljan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Get mypy to check the typing on the tests as well.

Would be nice with some typing on the ArrayApiIndexingAdapter too. But the dask version is unfortunately missing typing, although the pandas version has so maybe it's possible to draw inspiration from that one?

xarray/tests/test_array_api.py Outdated Show resolved Hide resolved
xarray/tests/test_array_api.py Outdated Show resolved Hide resolved
xarray/tests/test_array_api.py Outdated Show resolved Hide resolved
xarray/tests/test_array_api.py Outdated Show resolved Hide resolved
xarray/tests/test_array_api.py Outdated Show resolved Hide resolved
tomwhite and others added 6 commits July 19, 2022 13:44
Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
Copy link
Contributor

@dcherian dcherian left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@TomNicholas TomNicholas enabled auto-merge (squash) July 20, 2022 05:53
@TomNicholas TomNicholas merged commit 9f8d47c into pydata:main Jul 20, 2022
@jakirkham
Copy link

cc @rgommers (for awareness)

dcherian added a commit to keewis/xarray that referenced this pull request Jul 22, 2022
* main: (313 commits)
  Update whats-new
  Release notes for v2022.06.0 (pydata#6815)
  Drop multi-indexes when assigning to a multi-indexed variable (pydata#6798)
  Support NumPy array API (experimental) (pydata#6804)
  Add cumsum to DatasetGroupBy (pydata#6525)
  Refactor groupby binary ops code. (pydata#6789)
  Update DataArray.rename + docu (pydata#6665)
  Switch to T_DataArray and T_Dataset in concat (pydata#6784)
  Fix typos found by codespell (pydata#6794)
  Update groupby attrs tests (pydata#6787)
  Update map_blocks to use chunksizes property. (pydata#6776)
  Fix `DataArrayRolling.__iter__` with `center=True` (pydata#6744)
  [test-upstream] Update flox repo URL (pydata#6780)
  Move _infer_meta_data and _parse_size to utils (pydata#6779)
  Make the `sel` error more descriptive when `method` is unset (pydata#6774)
  Move Rolling tests to their own testing module (pydata#6777)
  [pre-commit.ci] pre-commit autoupdate (pydata#6773)
  move da and ds fixtures to conftest.py (pydata#6730)
  Bump EnricoMi/publish-unit-test-result-action from 1 to 2 (pydata#6770)
  Type shape methods (pydata#6767)
  ...
@tomwhite tomwhite mentioned this pull request Sep 22, 2022
1 task
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
topic-arrays related to flexible array support
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Use pytorch as backend for xarrays
6 participants