Generalized Ufuncs

EXPERIMENTAL FEATURE added to Version 0.18.0 and above - see disclaimer.

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, as well as provides a wrapper to make a Python function a generalized ufunc.

Usage

NumPy generalized ufunc

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.normal(size=(3, 10, 10), chunks=(2, 10, 10))

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

Wrap own Python function

gufunc can be used to make a Python function behave like a generalized ufunc:

x = da.random.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, but 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, **kwargs) Apply a generalized ufunc or similar python function to arrays.
as_gufunc([signature]) Decorator for dask.array.gufunc.
gufunc(pyfunc, **kwargs) Binds pyfunc into dask.array.apply_gufunc when called.