Generalized Ufuncs

NumPy provides the concept of generalized ufuncs. Generalized ufuncs are functions that distinguish the various dimensions of passed arrays in the two classes loop dimensions and core dimensions. To accomplish this, a signature is specified for NumPy generalized ufuncs.

Dask integrates interoperability with NumPy’s generalized ufuncs by adhering to respective ufunc protocol, and provides a wrapper to make a Python function a generalized ufunc.

Usage

NumPy Generalized UFuncs

Note

NumPy generalized ufuncs are currently (v1.14.3 and below) stored in inside np.linalg._umath_linalg and might change in the future.

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

Create Generalized UFuncs

It can be difficult to create your own GUFuncs without going into the CPython API. However, the Numba project does provide a nice implementation with their numba.guvectorize decorator. See Numba’s documentation for more information.

Wrap your own Python function

gufunc can be used to make a Python function behave like a generalized 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)

Instead of gufunc, also the as_gufunc decorator can be used for convenience:

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)

Disclaimer

This experimental generalized ufunc integration is not complete:

  • gufunc does not create a true generalized ufunc to be used with other input arrays besides Dask. I.e., at the moment, gufunc casts all input arguments to dask.array.Array

  • Inferring output_dtypes automatically is not implemented yet

API

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

Apply a generalized ufunc or similar python function to arrays.

as_gufunc([signature])

Decorator for dask.array.gufunc.

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

Binds pyfunc into dask.array.apply_gufunc when called.