Duck array ops try to import transpose from sparse

[brendan-m-murphy]

What happened?

Some of our tests that passed before the new year are failing now. The error is raised by some new xarray code, which might have first been on the Jan 2025 release of xarray.

This is the error:

    def transpose(array, axes=None):
        xp = get_array_namespace(array)
>       return xp.transpose(array, axes)
E       AttributeError: module 'sparse' has no attribute 'transpose'

I included the full traceback below. Since this is on our CI, I can't run xr.show_versions() directly, but I included the versions of xarray and sparse. I will try to reproduce it locally.

I don't have a MCVE. For context, the function that raises this error defines a ufunc that calls sparse.tensordot. I haven't had time to raise an issue for this, but about a year ago, I was having trouble using @ to multiply a sparse data array and a non-sparse data array, so I wrote a ufunc to do it. I think matrix mult. with one or more sparse data arrays works in some cases, and maybe this has been fixed in the interim.

This is the function that raises the error: https://github.com/openghg/openghg_inversions/blob/devel/openghg_inversions/array_ops.py#L75

Updates:

• the error from matrix multiplication with sparse matrices happens when the other data array is a dask array: I opened a separate issue for this: Matrix multiplication inconsistent for dask and sparse: dask @ sparse works, sparse @ dask fails #9934
• I reproduced the error locally, and updated the environment details below

What did you expect to happen?

Nothing has changed in the part of our code that is raising this error, so I would expect these tests to keep passing. The tests pass with sparse 15.4.0 and xarray 2024.11.0.

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

openghg_inversions/hbmcmc/hbmcmc.py:339: in fixedbasisMCMC
    fp_data = basis_functions_wrapper(
openghg_inversions/basis/_wrapper.py:118: in basis_functions_wrapper
    fp_data = fp_sensitivity(fp_all, basis_func=basis_data_array)
openghg_inversions/basis/_helpers.py:68: in fp_sensitivity
    sensitivity, site_bf = fp_sensitivity_single_site_basis_func(
openghg_inversions/basis/_helpers.py:159: in fp_sensitivity_single_site_basis_func
    sensitivity = sparse_xr_dot(basis_mat, H_all.fillna(0.0)).transpose("region", "time")
openghg_inversions/array_ops.py:170: in sparse_xr_dot
    return xr.apply_ufunc(func, da1, da2.as_numpy(), input_core_dims=input_core_dims, join="outer")
/opt/hostedtoolcache/Python/3.10.16/x64/lib/python3.10/site-packages/xarray/core/computation.py:1271: in apply_ufunc
    return apply_dataarray_vfunc(
/opt/hostedtoolcache/Python/3.10.16/x64/lib/python3.10/site-packages/xarray/core/computation.py:313: in apply_dataarray_vfunc
    result_var = func(*data_vars)
/opt/hostedtoolcache/Python/3.10.16/x64/lib/python3.10/site-packages/xarray/core/computation.py:733: in apply_variable_ufunc
    input_data = [
/opt/hostedtoolcache/Python/3.10.16/x64/lib/python3.10/site-packages/xarray/core/computation.py:735: in <listcomp>
    broadcast_compat_data(arg, broadcast_dims, core_dims)
/opt/hostedtoolcache/Python/3.10.16/x64/lib/python3.10/site-packages/xarray/core/computation.py:681: in broadcast_compat_data
    data = duck_array_ops.transpose(data, order)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

array = <COO: shape=(4, 293, 391), dtype=int64, nnz=114563, fill_value=0>
axes = (0, 2, 1)

    def transpose(array, axes=None):
        xp = get_array_namespace(array)
>       return xp.transpose(array, axes)
E       AttributeError: module 'sparse' has no attribute 'transpose'

Anything else we need to know?

As far as I can tell, sparse does not export transpose, so presumably any code that triggers the duck array transpose for a sparse array should raise this error. https://github.com/pydata/sparse/blob/0.15.4/sparse/__init__.py

Environment

INSTALLED VERSIONS ------------------ commit: None python: 3.10.14 (main, May 9 2024, 10:37:23) [Clang 15.0.0 (clang-1500.3.9.4)] python-bits: 64 OS: Darwin OS-release: 24.2.0 machine: arm64 processor: arm byteorder: little LC_ALL: None LANG: en_GB.UTF-8 LOCALE: ('en_GB', 'UTF-8') libhdf5: 1.12.2 libnetcdf: 4.9.3-development

xarray: 2025.1.0
pandas: 2.2.3
numpy: 1.26.4
scipy: 1.15.0
netCDF4: 1.6.5
pydap: None
h5netcdf: 1.4.1
h5py: 3.12.1
zarr: 2.18.3
cftime: 1.6.4.post1
nc_time_axis: 1.4.1
iris: None
bottleneck: None
dask: 2024.12.1
distributed: None
matplotlib: 3.10.0
cartopy: None
seaborn: None
numbagg: None
fsspec: 2024.12.0
cupy: None
pint: None
sparse: 0.15.4
flox: None
numpy_groupies: None
setuptools: 75.8.0
pip: 24.3.1
conda: None
pytest: 8.3.4
mypy: None
IPython: 8.31.0
sphinx: None

[keewis]

looks like the Array API defines the all-axis transpose as permute_dims with a tuple parameter (no default), so we might have to use that instead (and compute a default value)

[slevang]

Looks like this is used a grand total of one time in the codebase, so yeah we should probably switch to the proper Array API method permute_dims including renaming our duck_array_op. Would also be good to find a way to hit this code path in these tests, which ideally would have caught this.

[brendan-m-murphy]

If it helps for the tests, a ufunc that needs to permute a sparse array should trigger this, e.g.

da = xr.DataArray(COO.from_numpy(np.arange(12).reshape(3, 4)))
xr.apply_ufunc(np.sum, da, input_core_dims=[["dim_0"]], kwargs=dict(axis=0))