将 `xarray.apply_ufunc` 与 `np.linalg.pinv` 一起使用会返回 `dask.array` 的错误答案

【问题标题】：Using `xarray.apply_ufunc` with `np.linalg.pinv` returns an error with `dask.array`将 `xarray.apply_ufunc` 与 `np.linalg.pinv` 一起使用会返回 `dask.array` 的错误
【发布时间】：2023-02-04 01:39:09
【问题描述】：

运行以下 MWE 时出现错误：

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",
                        )

这给了我这个错误：

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'

即使直接使用dask.array.map_blocks，事情似乎开箱即用：

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

我在这里错过了什么？

【问题讨论】：

标签： python dask python-xarray

【解决方案1】：

错误消息似乎表明与 dask.array 一起使用时 np.linalg.pinv 的输出尺寸存在问题。使用 dask.array.map_blocks 时，它似乎工作正常。

问题可能是因为 xarray.apply_ufunc 目前与 dask.array 不兼容。要解决此问题，您可以考虑直接使用 dask.array.map_blocks 将函数 np.linalg.pinv 应用于您的数据。

【讨论】：