dask.array.apply_gufunc 具有不同形状的多个输出

Question

我正在尝试将 ufunc 应用于分块的可广播 dask 数组，这些数组产生多个不同形状的输出:

import dask.array as da # dask.__version__ is 1.2.0
import numpy as np

def func(A3, A2):
    return A3+A2, A2**2

A3 = da.from_array(np.random.randn(3,5,5), chunks=(3,2,2))
A2 = da.from_array(np.random.randn(  5,5), chunks=(  2,2))
ret = da.apply_gufunc(func, '(),()->(),()', A3, A2, output_dtypes=[float,float])

for r in ret:
    print(r)
    r.compute()

问题是 ret 中的两个输出假定形状为 (3,5,5)，然后在 .compute() 上失败与 ValueError: axes don't match array对于第二个输出，它应该是二维的。

我如何使用 apply_gufunc在这种情况下？

注:我可能宁愿使用 xarray.apply_ufunc在这种情况下，但不幸的是，它还没有多个输出可用(请参阅 here )。

Answer 1

这是一个有帮助的解决方法:

def func(A3, A2):
    return A3+A2, (A2**2)[np.newaxis,:]