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.

Sign up for GitHub

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

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

Merged
merged 6 commits into from
Sep 30, 2024
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
cc4f5bb
don't pass along `out` for `nanprod`
keewis Sep 22, 2024
a10e97a
add tests for `as_indexable`
keewis Sep 22, 2024
ecb47d6
allow using `__array_function__` as a fallback for missing array API …
keewis Sep 22, 2024
445a887
also check dask
keewis Sep 22, 2024
581f465
don't try to create a `dask` array if `dask` is not installed
keewis Sep 22, 2024
0a5ce3b
Merge branch 'main' into array-function-as-fallback
keewis Sep 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions xarray/core/indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -875,10 +875,10 @@ def as_indexable(array):
return PandasIndexingAdapter(array)
if is_duck_dask_array(array):
return DaskIndexingAdapter(array)
if hasattr(array, "__array_function__"):
return NdArrayLikeIndexingAdapter(array)
if hasattr(array, "__array_namespace__"):
return ArrayApiIndexingAdapter(array)
if hasattr(array, "__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.

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

Copy link
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
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.

return NdArrayLikeIndexingAdapter(array)

raise TypeError(f"Invalid array type: {type(array)}")

Expand Down
2 changes: 1 addition & 1 deletion xarray/core/nanops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,7 +162,7 @@ def nanstd(a, axis=None, dtype=None, out=None, ddof=0):

def nanprod(a, axis=None, dtype=None, out=None, min_count=None):
mask = isnull(a)
result = nputils.nanprod(a, axis=axis, dtype=dtype, out=out)
result = nputils.nanprod(a, axis=axis, dtype=dtype)
if min_count is not None:
return _maybe_null_out(result, axis, mask, min_count)
else:
Expand Down
68 changes: 68 additions & 0 deletions xarray/tests/test_indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -894,6 +894,74 @@ def test_posify_mask_subindexer(indices, expected) -> None:
np.testing.assert_array_equal(expected, actual)


class ArrayWithNamespace:
def __array_namespace__(self, version=None):
pass


class ArrayWithArrayFunction:
def __array_function__(self, func, types, args, kwargs):
pass


class ArrayWithNamespaceAndArrayFunction:
def __array_namespace__(self, version=None):
pass

def __array_function__(self, func, types, args, kwargs):
pass


def as_dask_array(arr, chunks):
try:
import dask.array as da
except ImportError:
return None

return da.from_array(arr, chunks=chunks)


@pytest.mark.parametrize(
["array", "expected_type"],
(
pytest.param(
indexing.CopyOnWriteArray(np.array([1, 2])),
indexing.CopyOnWriteArray,
id="ExplicitlyIndexed",
),
pytest.param(
np.array([1, 2]), indexing.NumpyIndexingAdapter, id="numpy.ndarray"
),
pytest.param(
pd.Index([1, 2]), indexing.PandasIndexingAdapter, id="pandas.Index"
),
pytest.param(
as_dask_array(np.array([1, 2]), chunks=(1,)),
indexing.DaskIndexingAdapter,
id="dask.array",
marks=requires_dask,
),
pytest.param(
ArrayWithNamespace(), indexing.ArrayApiIndexingAdapter, id="array_api"
),
pytest.param(
ArrayWithArrayFunction(),
indexing.NdArrayLikeIndexingAdapter,
id="array_like",
),
pytest.param(
ArrayWithNamespaceAndArrayFunction(),
indexing.ArrayApiIndexingAdapter,
id="array_api_with_fallback",
),
),
)
def test_as_indexable(array, expected_type):
actual = indexing.as_indexable(array)

assert isinstance(actual, expected_type)


def test_indexing_1d_object_array() -> None:
items = (np.arange(3), np.arange(6))
arr = DataArray(np.array(items, dtype=object))
Expand Down
Loading