广义Ufuncs

NumPy 提供了 广义ufuncs 的概念。广义ufuncs 是区分传递数组的各种维度的函数,这些维度分为循环维度和核心维度。为了实现这一点,为 NumPy 广义ufuncs 指定了 签名

Dask 通过遵循相应的 ufunc 协议,实现了与 NumPy 的通用 ufunc 的互操作性,并提供了一个包装器,使 Python 函数成为一个通用的 ufunc。

用法

NumPy 通用 UFuncs

备注

NumPy 的广义 ufuncs 当前(v1.14.3 及以下版本)存储在 np.linalg._umath_linalg 内部,未来可能会发生变化。

import dask.array as da
import numpy as np

x = da.random.default_rng().normal(size=(3, 10, 10), chunks=(2, 10, 10))

w, v = np.linalg._umath_linalg.eig(x, output_dtypes=(float, float))

创建广义的UFuncs

如果没有深入到 CPython API,创建自己的 GUFuncs 可能会很困难。然而,Numba 项目确实提供了一个很好的实现,即他们的 numba.guvectorize 装饰器。更多信息请参见 Numba 的文档

包装你自己的 Python 函数

gufunc 可以用来使一个 Python 函数表现得像一个广义的 ufunc:

x = da.random.default_rng().normal(size=(10, 5), chunks=(2, 5))

def foo(x):
    return np.mean(x, axis=-1)

gufoo = da.gufunc(foo, signature="(i)->()", output_dtypes=float, vectorize=True)

y = gufoo(x)

除了 gufunc,也可以使用 as_gufunc 装饰器以方便操作:

x = da.random.normal(size=(10, 5), chunks=(2, 5))

@da.as_gufunc(signature="(i)->()", output_dtypes=float, vectorize=True)
def gufoo(x):
    return np.mean(x, axis=-1)

y = gufoo(x)

免责声明

这个实验性的广义ufunc集成尚未完成:

  • gufunc 不会创建一个真正的通用ufunc,以便与其他输入数组(除了Dask)一起使用。也就是说,目前,gufunc 会将所有输入参数转换为 dask.array.Array

  • 自动推断 output_dtypes 尚未实现

API

apply_gufunc(func, signature, *args[, axes, ...])

将广义ufunc或类似的Python函数应用于数组。

as_gufunc([signature])

dask.array.gufunc 的装饰器。

gufunc(pyfunc, *[, signature, vectorize, ...])

在调用时将 pyfunc 绑定到 dask.array.apply_gufunc