Getting xarray.apply_ufunc and np.linalg.pinv to work with dask.array

[tomchor]

I'm having trouble parallelizing np.linalg.pinv using xarray.apply_ufunc. Here's a MWE that reproduces the error I'm getting:

import xarray as xr
import numpy as np
from numpy.linalg import pinv
import dask

data = np.random.randn(4, 4, 3, 2)

da = xr.DataArray(data=data, dims=("x", "y", "i", "j"),)

da = da.chunk(x=1, y=1)
da_inv = xr.apply_ufunc(pinv, da,
                        input_core_dims=[["i", "j"]],
                        output_core_dims=[["i", "j"]],
                        exclude_dims=set(("i", "j")),
                        dask = "parallelized",
                        )

This throws me this error:

Traceback (most recent call last):
  File "/glade/scratch/tomasc/tracer_inversion2/mwe.py", line 14, in <module>
    da_inv = xr.apply_ufunc(pinv, da,
  File "/glade/u/home/tomasc/miniconda3/envs/py310/lib/python3.10/site-packages/xarray/core/computation.py", line 1204, in apply_ufunc
    return apply_dataarray_vfunc(
  File "/glade/u/home/tomasc/miniconda3/envs/py310/lib/python3.10/site-packages/xarray/core/computation.py", line 315, in apply_dataarray_vfunc
    result_var = func(*data_vars)
  File "/glade/u/home/tomasc/miniconda3/envs/py310/lib/python3.10/site-packages/xarray/core/computation.py", line 771, in apply_variable_ufunc
    result_data = func(*input_data)
  File "/glade/u/home/tomasc/miniconda3/envs/py310/lib/python3.10/site-packages/xarray/core/computation.py", line 747, in func
    res = da.apply_gufunc(
  File "/glade/u/home/tomasc/miniconda3/envs/py310/lib/python3.10/site-packages/dask/array/gufunc.py", line 489, in apply_gufunc
    core_output_shape = tuple(core_shapes[d] for d in ocd)
  File "/glade/u/home/tomasc/miniconda3/envs/py310/lib/python3.10/site-packages/dask/array/gufunc.py", line 489, in <genexpr>
    core_output_shape = tuple(core_shapes[d] for d in ocd)
KeyError: 'dim0'

When using dask.array.map_blocks directly, things seem to work right out of the box:

data_inv = dask.array.map_blocks(pinv, da.data).compute() # works!

So I guess I must be missing something simple here?

Thanks!

[TomNicholas]

You were almost there, you just needed to add the sizes of new the output dimensions by including the kwarg

dask_gufunc_kwargs={'output_sizes': {'i': 2, 'j': 3}}

It does sort of say this in the docstring for apply_ufunc but it could definitely be clearer!

That's a very unhelpful error, but it's ultimately being thrown because the keys 'i' and 'j' don't exist in the dict of expected sizes of the output (because you didn't provide them). We should improve that error!

[TomNicholas]

It is also possible to use (the somewhat experimental) xarray.map_blocks for this, but you first have to wrap np.linalg.pinv to deal with the fact it won't return an xarray object:

def xr_pinv(da: xr.DataArray) -> xr.DataArray:
    new_data = np.linalg.pinv(da.data)
    return xr.DataArray(new_data, dims=['x', 'y', 'j', 'i'])

xr.map_blocks(xr_pinv, da)

EDIT: I don't think I would recommend this - you are basically silently applying the pinv function to only the last two dimensions of your DataArray, which could change under a transposition without you noticing. apply_ufunc with two core dimensions as you did is the correct way to wrap pinv.

[TomNicholas]

Oh huh the terrible error message is actually a result of you choosing the same names for the new output dimensions as the input dimensions. If you choose to name the outputs as e.g. 'a', 'b' instead you get a much clearer message:

da_inv = xr.apply_ufunc(
    pinv, da,
    input_core_dims=[["i", "j"]],
    output_core_dims=[["a", "b"]],
    exclude_dims=set(("i", "j")),
    dask = "parallelized",
)

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[34], line 1
----> 1 da_inv = xr.apply_ufunc(pinv, da,
      2                         input_core_dims=[["i", "j"]],
      3                         output_core_dims=[["a", "b"]],
      4                         exclude_dims=set(("i", "j")),
      5                         dask = "parallelized",
      6                         #dask_gufunc_kwargs={'output_sizes': {'a': 2, 'b': 3}},
      7                         )

File ~/Documents/Work/Code/xarray/xarray/core/computation.py:1194, in apply_ufunc(func, input_core_dims, output_core_dims, exclude_dims, vectorize, join, dataset_join, dataset_fill_value, keep_attrs, kwargs, dask, output_dtypes, output_sizes, meta, dask_gufunc_kwargs, *args)
   1192 # feed DataArray apply_variable_ufunc through apply_dataarray_vfunc
   1193 elif any(isinstance(a, DataArray) for a in args):
-> 1194     return apply_dataarray_vfunc(
   1195         variables_vfunc,
   1196         *args,
   1197         signature=signature,
   1198         join=join,
   1199         exclude_dims=exclude_dims,
   1200         keep_attrs=keep_attrs,
   1201     )
   1202 # feed Variables directly through apply_variable_ufunc
   1203 elif any(isinstance(a, Variable) for a in args):

File ~/Documents/Work/Code/xarray/xarray/core/computation.py:303, in apply_dataarray_vfunc(func, signature, join, exclude_dims, keep_attrs, *args)
    298 result_coords, result_indexes = build_output_coords_and_indexes(
    299     args, signature, exclude_dims, combine_attrs=keep_attrs
    300 )
    302 data_vars = [getattr(a, "variable", a) for a in args]
--> 303 result_var = func(*data_vars)
    305 out: tuple[DataArray, ...] | DataArray
    306 if signature.num_outputs > 1:

File ~/Documents/Work/Code/xarray/xarray/core/computation.py:727, in apply_variable_ufunc(func, signature, exclude_dims, dask, output_dtypes, vectorize, keep_attrs, dask_gufunc_kwargs, *args)
    725 for key in signature.all_output_core_dims:
    726     if key not in signature.all_input_core_dims and key not in output_sizes:
--> 727         raise ValueError(
    728             f"dimension '{key}' in 'output_core_dims' needs corresponding (dim, size) in 'output_sizes'"
    729         )
    731 def func(*arrays):
    732     import dask.array as da

ValueError: dimension 'b' in 'output_core_dims' needs corresponding (dim, size) in 'output_sizes'

Perhaps we should forbid this? Or at least catch the error sneaking through...

[TomNicholas]

The error should be fixed to be clearer in #7509

[tomchor]

You were almost there, you just needed to add the sizes of new the output dimensions by including the kwarg

dask_gufunc_kwargs={'output_sizes': {'i': 2, 'j': 3}}

Ah, so close! 😂

Thanks for the detailed answer (and for improving the error message in #7509). I learned a lot! Sorry I missed that from the docs 😬