广义 Ufuncs

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

Dask 通过遵循相应的 ufunc 协议,集成了与 NumPy 广义 ufuncs 的互操作性,并提供了一个包装器来使 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