Skip to content

allow using __array_function__ as a fallback for missing Array API functions - #9530

Merged
TomNicholas merged 6 commits into
pydata:mainfrom
keewis:array-function-as-fallback
Sep 30, 2024
Merged

TomNicholas merged 6 commits into
pydata:mainfrom
keewis:array-function-as-fallback

Conversation

@keewis

@keewis keewis commented Sep 22, 2024 •

Copy link
Copy Markdown
Collaborator

As discussed in #8834 (comment), this swaps the order of special casing for types implementing both __array_function__ and __array_namespace__, such that the Array API is preferred over __array_function__.

I've also disabled passing the out kwarg to nanprod, as that was the only occurrence where we actually do that, all other nan* functions silently ignore that kwarg.

  • Tests added
  • User visible changes (including notable bug fixes) are documented in whats-new.rst

cc @tomwhite

@TomNicholas TomNicholas added topic-arrays related to flexible array support array API standard Support for the Python array API standard labels Sep 22, 2024
@TomNicholas

Copy link
Copy Markdown
Member

This seems good, but does it cover this comment? #8834 (comment)

@tomwhite

Copy link
Copy Markdown
Contributor

Thanks @keewis - this would be very helpful!

@keewis

keewis commented Sep 23, 2024

Copy link
Copy Markdown
Collaborator Author

does it cover this comment? #8834 (comment)

I don't think it does, but I think array types that implement both __array_{ufunc,function}__ and __array_namespace__ would prefer to use __array_namespace__, so I believe we don't change anything for that (not that I know of any array type that actually fully implements both besides numpy).

@dcherian

dcherian commented Sep 23, 2024 •

Copy link
Copy Markdown
Contributor

This almost certainly breaks vectorized indexing with cupy if you don't implement the fallback vindex using explicit_indexing_adapter.

Though how are the pint tests passing? Do we have indexing tests with pint?

@keewis

keewis commented Sep 23, 2024 •

Copy link
Copy Markdown
Collaborator Author

This almost certainly breaks vectorized indexing with cupy

As far as I can tell, cupy does not implement __array_namespace__ (unless you used cupy.array_api or array_api_compat.cupy to create the array, where the former apparently is sunset in cupy=14), so it will still use the NdArrayLikeIndexingAdapter. Same with pint, it doesn't implement __array_namespace__.

That said, we should definitely try to provide an implementation for vindex, I just don't think it is as urgent.

@dcherian

Copy link
Copy Markdown
Contributor

OK yes, lot less urgent then. I thought they had all implemented __array_namespace__ for their main api like numpy.

@keewis

keewis commented Sep 23, 2024

Copy link
Copy Markdown
Collaborator Author

cupy plans to do so in v14 (see cupy/cupy#8470), they try to be a drop-in replacement for numpy – though obviously there is some lag.

Comment thread xarray/core/indexing.py
return NdArrayLikeIndexingAdapter(array)
if hasattr(array, "__array_namespace__"):
return ArrayApiIndexingAdapter(array)
if hasattr(array, "__array_function__"):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

I'm confused now, why does the indexing adapter affect what's used for reductions?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

I don't think that's the case. What I mentioned in the meeting was that we could use __array_function__ as a fallback for nan* reductions, whereas choosing the wrong indexing adapter will cause xarray to try indexing with a method the array type might not support.

I guess this is kinda hard to figure out implicitly, and I'd love to have a way (a special attribute like __xarray_preferred_array_protocol__?) to explicitly communicate which protocol is preferred. That way, we won't have to worry about cupy and pint implementing the Array API, and cubed can just set that attribute to state that it considers __array_function__ a fallback and instead prefers to use __array_namespace__ indexing.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

array_function as a fallback for nan* reductions,

👍 sorry I got confused.

@dcherian dcherian added the plan to merge Final call for comments label Sep 25, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

array API standard Support for the Python array API standard plan to merge Final call for comments topic-arrays related to flexible array support

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants