Namespace-aware `xarray.ufuncs` #9776

slevang · 2024-11-13T04:49:14Z

Closes Compatibility with the Array API standard #7848 (partially)
Tests added
User visible changes (including notable bug fixes) are documented in whats-new.rst
New functions/methods are listed in api.rst

Re-implement the old xarray.ufuncs module to allow generic ufunc handling for array types that don't implement __array_ufunc__:

import jax.numpy as jnp
import numpy as np
import xarray as xr
import xarray.ufuncs as xu

x = xr.DataArray(jnp.asarray([1, 2, 3]))
print(type(xu.sin(x).data))
print(type(np.sin(x).data))

# <class 'jaxlib.xla_extension.ArrayImpl'>
# <class 'numpy.ndarray'>

slevang · 2024-11-13T14:10:29Z

xarray/ufuncs.py

+    elif isinstance(obj, array_type("dask")):
+        _walk_array_namespaces(obj._meta, namespaces)
+    else:
+        namespace = getattr(obj, "__array_namespace__", None)


Should we ever prioritize dispatching with np.func via __array_ufunc__ (if it exists) over the library's __array_namespace__().func?

__array_ufunc__ is a more generic protocol, intended to support arbitrary new ufuncs without requiring an array library to be aware of them.

In practice there are very few examples of ufuncs defined outside of NumPy itself, and we wouldn't need to support them here because we are explicitly listing supported ufuncs in this module. I guess the one example that comes to mind would be the rare cases where NumPy adds a new ufunc and an array wrapping library like Dask hasn't written a wrapper yet.

At this point, I think going "all in" on __array_namespace__ is the right call.

I guess my other comment would be the main reason to consider __array_ufunc__. Some duck arrays don't implement all ufuncs. So either of these approaches would solve the same problem.

slevang · 2024-11-13T14:13:21Z

xarray/ufuncs.py

+            )
+            func = getattr(np, self._name)
+
+        return xr.apply_ufunc(func, *args, dask="parallelized", **kwargs)


Is there ever a reason to use dask's ufuncs with dask="allowed" instead of the appropriate _meta array's namespace and dask="parallelized"? With jax for example, which doesn't have __array_ufunc__, this ends up converting to numpy. So it would have to be special cased.

In user code using xr.apply_ufunc there is - dask='allowed' can be used to rechunk along a core dimension e.g. by applying a dask reduction ufunc along that dimension. Not sure if that's relevant here though.

These are all elementwise so no core dimensions

Is there ever a reason to use dask's ufuncs with dask="allowed" instead of the appropriate _meta array's namespace and dask="parallelized"?

Yes, this feels like a cleaner solution to me.

With jax for example, which doesn't have array_ufunc, this ends up converting to numpy. So it would have to be special cased

Is the concern here Dask wrapping JAX?

Generally, I think it's best for xarray to avoid introspecting into wrapped array types in Xarray, and leave nested wrapping up to other libraries, which can better understand their own implementation details. So Dask wrapping JAX should be fixed in Dask.

So Dask wrapping JAX should be fixed in Dask.

Totally fair. It looks like basically the same effort here would be required in dask then, because dask's ufuncs are all simple wrappers around the numpy version so they aren't aware of the namespace.

slevang · 2024-11-13T14:15:08Z

xarray/ufuncs.py

+
+
+# Auto generate from the public numpy ufuncs
+np_ufuncs = {name for name in dir(np) if isinstance(getattr(np, name), np.ufunc)}


Can hard code these if preferred?

Yes, I would suggest hard coding these if possible, ideally as something like:

sin = _unary_ufunc('sin') add = _binary_ufunc('add') ...

The reason to prefer this is that static typing does not evaluate loops or dynamically defined functions. So otherwise xarray.ufuncs.sin() will not be recognized as valid by tools like mypy.

Ideally (could be done in a follow-up PR), these functions could be annotated to follow the appropriate type casting rules, so type checkers would know that xarray.ufuncs.sin(dataset) returns another Dataset.

In some cases, we use a script to generate all these special methods. I don't think that should be necessary here, but it still may be worth a look to understand the type casting rules:
https://github.com/pydata/xarray/blob/d5f84dd1ef4c023cf2ea0a38866c9d9cd50487e7/xarray/util/generate_ops.py
https://github.com/pydata/xarray/blob/d5f84dd1ef4c023cf2ea0a38866c9d9cd50487e7/xarray/core/_typed_ops.py

slevang · 2024-11-13T14:15:48Z

xarray/ufuncs.py

+
+# Auto generate from the public numpy ufuncs
+np_ufuncs = {name for name in dir(np) if isinstance(getattr(np, name), np.ufunc)}
+excluded_ufuncs = {"divmod", "frexp", "isnat", "matmul", "modf", "vecdot"}


These are the ones that didn't immediately work. There are also other ufunc like things that aren't technically np.ufunc subclasses that we could add. I saw angle and iscomplex were special cased before.

Maybe worth noting that the reason why matmul and vecdot doesn't work is that they are "generalized ufuncs" that use core dimensions.

divmod, frexp and modf doesn't work because they return multiple arrays.

I'm not sure why isnat didn't work for you. Did you test it with datetime dtypes?

Did you test it with datetime dtypes?

No, this was just a real quick initial pass. Will add this with a special test case.

Do you have an opinion about adding any of the ones with multiple return values? Seems low priority to me.

Same question for the odd balls like angle, iscomplex, isreal, etc?

shoyer · 2024-11-16T21:23:19Z

xarray/ufuncs.py

+        if func is None:
+            warnings.warn(
+                f"Function {self._name} not found in {xp.__name__}, falling back to numpy",
+                stacklevel=2,
+            )
+            func = getattr(np, self._name)


I would lean towards skipping this fall-back, unless there are particularly motivating cases.

Motivation here would be duck arrays that implement __array_ufunc__ and don't implement the full suite of numpy ufuncs. I ran into this with sparse. Not sure the full delta list, but I see they don't have sin/cos for example. In this case, np.cos(x_sparse) works but xp.cos(x_sparse) fails, which is a little weird. Not the most elegant solution though, I agree.

I think it's intentional than xp.cos(x_sparse) fails, because cos(0) != 0, so the result is no longer sparse.

Actually I was wrong, there is a sparse.cos and a bunch of others, although they don't appear in the API docs. It seems sparse's general approach to these is compute elementwise on the valid data, and then modify the fill_value as required for the empty regions.

There are still 40-some functions that fail without this fallback, although generally more niche. With the fallback, all work and output a sparse array (no auto densification):

absolute, arccos, arccosh, arcsin, arcsinh, arctan, arctan2, arctanh, bitwise_count, cbrt, conj, conjugate, copysign, deg2rad, degrees, exp2, expm1x, fabs, float_power, fmax, fmin, fmod, gcd, heaviside, hypot, invert, isreal, lcm, ldexp, left_shift, logaddexp2, maximum, minimum, mod, nextafter, power, rad2deg, radians, reciprocal, right_shift, rint, signbit, spacing, true_divide

shoyer · 2024-11-16T21:24:32Z

xarray/ufuncs.py

+    elif isinstance(obj, array_type("dask")):
+        _walk_array_namespaces(obj._meta, namespaces)


Why not sure Dask's __array_namespace__ instead? That feels a little cleaner than special case logic for dask.array.

shoyer · 2024-11-16T21:26:32Z

xarray/ufuncs.py

+            )
+            func = getattr(np, self._name)
+
+        return xr.apply_ufunc(func, *args, dask="parallelized", **kwargs)


Is there ever a reason to use dask's ufuncs with dask="allowed" instead of the appropriate _meta array's namespace and dask="parallelized"?

Yes, this feels like a cleaner solution to me.

With jax for example, which doesn't have array_ufunc, this ends up converting to numpy. So it would have to be special cased

Is the concern here Dask wrapping JAX?

Generally, I think it's best for xarray to avoid introspecting into wrapped array types in Xarray, and leave nested wrapping up to other libraries, which can better understand their own implementation details. So Dask wrapping JAX should be fixed in Dask.

shoyer · 2024-11-16T21:37:56Z

xarray/ufuncs.py

+
+
+# Auto generate from the public numpy ufuncs
+np_ufuncs = {name for name in dir(np) if isinstance(getattr(np, name), np.ufunc)}


Yes, I would suggest hard coding these if possible, ideally as something like:

sin = _unary_ufunc('sin') add = _binary_ufunc('add') ...

The reason to prefer this is that static typing does not evaluate loops or dynamically defined functions. So otherwise xarray.ufuncs.sin() will not be recognized as valid by tools like mypy.

Ideally (could be done in a follow-up PR), these functions could be annotated to follow the appropriate type casting rules, so type checkers would know that xarray.ufuncs.sin(dataset) returns another Dataset.

In some cases, we use a script to generate all these special methods. I don't think that should be necessary here, but it still may be worth a look to understand the type casting rules:
https://github.com/pydata/xarray/blob/d5f84dd1ef4c023cf2ea0a38866c9d9cd50487e7/xarray/util/generate_ops.py
https://github.com/pydata/xarray/blob/d5f84dd1ef4c023cf2ea0a38866c9d9cd50487e7/xarray/core/_typed_ops.py

xarray/ufuncs.py

shoyer · 2024-11-16T21:46:54Z

xarray/ufuncs.py

+    elif isinstance(obj, array_type("dask")):
+        _walk_array_namespaces(obj._meta, namespaces)
+    else:
+        namespace = getattr(obj, "__array_namespace__", None)


__array_ufunc__ is a more generic protocol, intended to support arbitrary new ufuncs without requiring an array library to be aware of them.

In practice there are very few examples of ufuncs defined outside of NumPy itself, and we wouldn't need to support them here because we are explicitly listing supported ufuncs in this module. I guess the one example that comes to mind would be the rare cases where NumPy adds a new ufunc and an array wrapping library like Dask hasn't written a wrapper yet.

At this point, I think going "all in" on __array_namespace__ is the right call.

shoyer · 2024-11-16T21:51:31Z

xarray/ufuncs.py

+
+# Auto generate from the public numpy ufuncs
+np_ufuncs = {name for name in dir(np) if isinstance(getattr(np, name), np.ufunc)}
+excluded_ufuncs = {"divmod", "frexp", "isnat", "matmul", "modf", "vecdot"}


Maybe worth noting that the reason why matmul and vecdot doesn't work is that they are "generalized ufuncs" that use core dimensions.

divmod, frexp and modf doesn't work because they return multiple arrays.

I'm not sure why isnat didn't work for you. Did you test it with datetime dtypes?

…to array-namespace-ufuncs

slevang · 2024-11-17T14:37:29Z

xarray/ufuncs.py

+    def __init__(self, name):
+        self._name = name
+
+    def __call__(self, x, **kwargs):


@shoyer lmk if this is what you had in mind for a separating unary from binary funcs. I assume this is for typing purposes.

Close! I added a more detailed note below.

shoyer · 2024-11-17T20:48:02Z

xarray/ufuncs.py

+def _create_op(name):
+    if not hasattr(np, name):
+        # handle older numpy versions with missing array api standard aliases
+        if np.lib.NumpyVersion(np.__version__) < "2.0.0":
+            return _UnavailableUfunc(name)
+        raise ValueError(f"'{name}' is not a valid numpy function")
+
+    np_func = getattr(np, name)
+    if hasattr(np_func, "nin") and np_func.nin == 2:
+        func = _BinaryUfunc(name)
+    else:
+        func = _UnaryUfunc(name)
+
+    func.__name__ = name
+    doc = getattr(np, name).__doc__
+
+    doc = _remove_unused_reference_labels(_skip_signature(_dedent(doc), name))
+
+    func.__doc__ = (
+        f"xarray specific variant of numpy.{name}. Handles "
+        "xarray objects by dispatching to the appropriate "
+        "function for the underlying array type.\n\n"
+        f"Documentation from numpy:\n\n{doc}"
+    )
+    return func


Type checkers can't evaluate code at runtime, so you really want to write something that prescribes the type signature statically, like:

abs = _UnaryUfunc('abs')

or

def _unary_ufunc(name: str) -> _UnaryUfunc: func = _UnaryUfunc(name) func.__doc__ = ... return func abs = _unary_ufunc('abs')

When you write abs = _create_op('abs'), type checkers think the type of abs could be any of _UnavailableUfunc or _BinaryUfunc or _UnaryUfunc. (You can check this with reveal_type(abs) if you're curious.)

xarray/ufuncs.py

shoyer · 2024-11-17T20:50:26Z

xarray/ufuncs.py

+    def __init__(self, name):
+        self._name = name
+
+    def __call__(self, x, **kwargs):


Close! I added a more detailed note below.

xarray/ufuncs.py

shoyer · 2024-11-17T20:51:46Z

This looks great! Could you please also add a brief note to whats-new.rst?

slevang · 2024-11-17T21:21:57Z

This looks great! Could you please also add a brief note to whats-new.rst?

Yes will do. Shall we add these back to the main api docs as well, same as before minus the deprecation notice?

shoyer · 2024-11-18T01:41:45Z

Yes, please!

…

On Sun, Nov 17, 2024 at 1:22 PM Sam Levang ***@***.***> wrote: This looks great! Could you please also add a brief note to whats-new.rst? Yes will do. Shall we add these back to the main api docs as well, same as before <https://github.com/pydata/xarray/pull/6491/files#diff-9ceb22fb61ab27613e143e733f0d2ef128c1f72125ae3c6324465955bd093f93> minus the deprecation notice? — Reply to this email directly, view it on GitHub <#9776 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVVI36XVHPN3ZAPCCSL2BECIVAVCNFSM6AAAAABRVS4J72VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDIOBRGU4DEOJUGM> . You are receiving this because you were mentioned.Message ID: ***@***.***>

xarray/ufuncs.py

* main: (24 commits) Bump minimum versions (#9796) Namespace-aware `xarray.ufuncs` (#9776) Add prettier and pygrep hooks to pre-commit hooks (#9644) `rolling.construct`: Add `sliding_window_kwargs` to pipe arguments down to `sliding_window_view` (#9720) Bump codecov/codecov-action from 4.6.0 to 5.0.2 in the actions group (#9793) Buffer types (#9787) Add download stats badges (#9786) Fix open_mfdataset for list of fsspec files (#9785) add 'User-Agent'-header to pooch.retrieve (#9782) Optimize `ffill`, `bfill` with dask when `limit` is specified (#9771) fix cf decoding of grid_mapping (#9765) Allow wrapping `np.ndarray` subclasses (#9760) Optimize polyfit (#9766) Use `map_overlap` for rolling reductions with Dask (#9770) fix html repr indexes section (#9768) Bump pypa/gh-action-pypi-publish from 1.11.0 to 1.12.2 in the actions group (#9763) unpin array-api-strict, as issues are resolved upstream (#9762) rewrite the `min_deps_check` script (#9754) CI runs ruff instead of pep8speaks (#9759) Specify copyright holders in main license file (#9756) ...

* main: Bump minimum versions (pydata#9796) Namespace-aware `xarray.ufuncs` (pydata#9776) Add prettier and pygrep hooks to pre-commit hooks (pydata#9644) `rolling.construct`: Add `sliding_window_kwargs` to pipe arguments down to `sliding_window_view` (pydata#9720) Bump codecov/codecov-action from 4.6.0 to 5.0.2 in the actions group (pydata#9793) Buffer types (pydata#9787) Add download stats badges (pydata#9786) Fix open_mfdataset for list of fsspec files (pydata#9785) add 'User-Agent'-header to pooch.retrieve (pydata#9782) Optimize `ffill`, `bfill` with dask when `limit` is specified (pydata#9771) fix cf decoding of grid_mapping (pydata#9765) Allow wrapping `np.ndarray` subclasses (pydata#9760) Optimize polyfit (pydata#9766) Use `map_overlap` for rolling reductions with Dask (pydata#9770)

initial namespace-aware implementation

47610a6

slevang commented Nov 13, 2024

View reviewed changes

slevang marked this pull request as ready for review November 13, 2024 14:17

slevang mentioned this pull request Nov 13, 2024

Compatibility with the Array API standard #7848

Open

slevang added 2 commits November 13, 2024 11:29

Merge branch 'main' into array-namespace-ufuncs

596ee93

use np subclass, test duck dask arrays

6d9bfbe

TomNicholas added topic-arrays related to flexible array support array API standard Support for the Python array API standard labels Nov 14, 2024

Merge branch 'main' into array-namespace-ufuncs

a01085d

dcherian requested a review from keewis November 15, 2024 16:23

shoyer reviewed Nov 16, 2024

View reviewed changes

slevang added 5 commits November 16, 2024 22:22

remove dask special casing and numpy fallback

3799b4c

add isnat

0f929bd

hard code the supported ufuncs

b7c58cf

Merge branch 'array-namespace-ufuncs' of github.com:slevang/xarray in…

bba3183

…to array-namespace-ufuncs

handle np versions, separate unary/binary path

6a55ea6

slevang commented Nov 17, 2024

View reviewed changes

shoyer reviewed Nov 17, 2024

View reviewed changes

dcherian mentioned this pull request Nov 18, 2024

release notes for v2024.11.0 #9794

Merged

6 tasks

slevang added 3 commits November 18, 2024 11:47

explicit unary/binary creators

900ca17

add to api docs

2da600d

add whats new

22e678c

shoyer reviewed Nov 18, 2024

View reviewed changes

xarray/ufuncs.py Outdated Show resolved Hide resolved

slevang added 3 commits November 18, 2024 13:42

move numpy version check to tests

21dd468

fix docs for aliased np funcs

127b591

fix whats new

81d6c07

shoyer approved these changes Nov 18, 2024

View reviewed changes

Merge branch 'main' into array-namespace-ufuncs

06d3051

dcherian merged commit 5d70f4d into pydata:main Nov 18, 2024
29 checks passed

slevang mentioned this pull request Nov 18, 2024

Improved duck array wrapping #9798

Merged

4 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Namespace-aware `xarray.ufuncs` #9776

Namespace-aware `xarray.ufuncs` #9776

slevang commented Nov 13, 2024 •

edited

Loading

slevang Nov 13, 2024

shoyer Nov 16, 2024

slevang Nov 17, 2024

slevang Nov 13, 2024

TomNicholas Nov 14, 2024

slevang Nov 14, 2024

shoyer Nov 16, 2024

slevang Nov 17, 2024

slevang Nov 13, 2024

shoyer Nov 16, 2024

slevang Nov 13, 2024

shoyer Nov 16, 2024

slevang Nov 17, 2024

shoyer Nov 16, 2024

slevang Nov 17, 2024

shoyer Nov 17, 2024

slevang Nov 17, 2024

shoyer Nov 16, 2024

shoyer Nov 16, 2024

shoyer Nov 16, 2024

shoyer Nov 16, 2024

shoyer Nov 16, 2024

slevang Nov 17, 2024

shoyer Nov 17, 2024

shoyer Nov 17, 2024

shoyer Nov 17, 2024

shoyer commented Nov 17, 2024

slevang commented Nov 17, 2024

shoyer commented Nov 18, 2024 via email



		# Auto generate from the public numpy ufuncs
		np_ufuncs = {name for name in dir(np) if isinstance(getattr(np, name), np.ufunc)}

		elif isinstance(obj, array_type("dask")):
		_walk_array_namespaces(obj._meta, namespaces)

Namespace-aware xarray.ufuncs #9776

Namespace-aware xarray.ufuncs #9776

Conversation

slevang commented Nov 13, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shoyer commented Nov 17, 2024

slevang commented Nov 17, 2024

shoyer commented Nov 18, 2024 via email

Namespace-aware `xarray.ufuncs` #9776

Namespace-aware `xarray.ufuncs` #9776

slevang commented Nov 13, 2024 •

edited

Loading