-
-
Notifications
You must be signed in to change notification settings - Fork 1.1k
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
Conversation
xarray/core/variable.py
Outdated
if isinstance(data, NON_NUMPY_SUPPORTED_ARRAY_TYPES): | ||
if ( | ||
isinstance(data, NON_NUMPY_SUPPORTED_ARRAY_TYPES) | ||
or hasattr(data, "__array_function__") |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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?
xarray/core/indexing.py
Outdated
value = value[(slice(None),) * axis + (subkey, Ellipsis)] | ||
return value | ||
else: | ||
assert isinstance(key, VectorizedIndexer) |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Replaced with a TypeError
xarray/core/indexing.py
Outdated
if isinstance(key, BasicIndexer): | ||
self.array[key.tuple] = value | ||
elif isinstance(key, OuterIndexer): | ||
self.array[key.tuple] = value |
There was a problem hiding this comment.
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?
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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Correct - fixed.
Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
Co-authored-by: Illviljan <14371165+Illviljan@users.noreply.github.com>
There was a problem hiding this 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?
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>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM!
cc @rgommers (for awareness) |
* 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) ...
whats-new.rst
api.rst
Note: I haven't actually tested this with pytorch (which is the motivating example for #3232).