Failure on `main` testing xarray.cross

[TomNicholas]

What happened?

Failure on main testing xarray.cross

https://github.com/pydata/xarray/actions/runs/10298827771/job/28504816043#step:10:838

Possibly either a bug in apply_ufunc or an incorrect use of it.

What did you expect to happen?

No response

Minimal Complete Verifiable Example

No response

MVCE confirmation

• Minimal example — the example is as focused as reasonably possible to demonstrate the underlying issue in xarray.
• Complete example — the example is self-contained, including all data and the text of any traceback.
• Verifiable example — the example copy & pastes into an IPython prompt or Binder notebook, returning the result.
• New issue — a search of GitHub Issues suggests this is not a duplicate.
• Recent environment — the issue occurs with the latest version of xarray and its dependencies.

Relevant log output

No response

Anything else we need to know?

No response

Environment

main

[kmuehlbauer]

It looks like this surfaced with new numpy 2.0.1 release in connection with dask 2024.8.0.

• breaks with dask 2024.8.0 and numpy 2.0.1
• works with dask 2024.8.0 and numpy 2.0.0
• works with dask 2024.7.1 and numpy 2.0.1

[kmuehlbauer]

So from what I can tell the align-path in below code should align a and b. But it returns multiple chunks instead of one chunk.

In the following elif-path there is special handling to rechunk. Hope this helps to figure out a solution.

xarray/xarray/core/computation.py

Lines 1669 to 1705 in b22c429

	if a.sizes[dim] != b.sizes[dim]:
	# Arrays have different sizes. Append zeros where the smaller
	# array is missing a value, zeros will not affect np.cross:
	if (
	not isinstance(a, Variable) # Only used to make mypy happy.
	and dim in getattr(a, "coords", {})
	and not isinstance(b, Variable) # Only used to make mypy happy.
	and dim in getattr(b, "coords", {})
	):
	# If the arrays have coords we know which indexes to fill
	# with zeros:
	a, b = align(
	a,
	b,
	fill_value=0,
	join="outer",
	exclude=set(all_dims) - {dim},
	)
	elif min(a.sizes[dim], b.sizes[dim]) == 2:
	# If the array doesn't have coords we can only infer
	# that it has composite values if the size is at least 2.
	# Once padded, rechunk the padded array because apply_ufunc
	# requires core dimensions not to be chunked:
	if a.sizes[dim] < b.sizes[dim]:
	a = a.pad({dim: (0, 1)}, constant_values=0)
	# TODO: Should pad or apply_ufunc handle correct chunking?
	a = a.chunk({dim: -1}) if is_duck_dask_array(a.data) else a
	else:
	b = b.pad({dim: (0, 1)}, constant_values=0)
	# TODO: Should pad or apply_ufunc handle correct chunking?
	b = b.chunk({dim: -1}) if is_duck_dask_array(b.data) else b
	else:
	raise ValueError(
	f"{dim!r} on {'a' if a.sizes[dim] == 1 else 'b'} is incompatible:"
	" dimensions without coordinates must have have a length of 2 or 3"
	)