Generalized Ufuncs

NumPy <https://www.numpy.org>_ provides the concept of generalized ufuncs <https://docs.scipy.org/doc/numpy/reference/c-api/generalized-ufuncs.html>. 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 <https://docs.scipy.org/doc/numpy/reference/c-api/generalized-ufuncs.html#details-of-signature> is specified for NumPy generalized ufuncs.

Dask <https://dask.org/>_ integrates interoperability with NumPy's generalized ufuncs by adhering to respective ufunc protocol <https://docs.scipy.org/doc/numpy/reference/arrays.classes.html#numpy.class.__array_ufunc__>_, and provides a wrapper to make a Python function a generalized ufunc.

Usage

NumPy Generalized UFuncs

.. note::

    `NumPy <https://www.numpy.org>`_ generalized ufuncs are currently (v1.14.3 and below) stored in
    inside ``np.linalg._umath_linalg`` and might change in the future.


.. code-block:: python

    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 <https://numba.pydata.org>_ project does provide a nice implementation with their numba.guvectorize decorator. See Numba's documentation <https://numba.pydata.org/numba-doc/dev/user/vectorize.html#the-guvectorize-decorator>_ for more information.

Wrap your own Python function

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


.. code-block:: python

    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:


.. code-block:: python

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

.. currentmodule:: dask.array.gufunc

.. autosummary::
   apply_gufunc
   as_gufunc
   gufunc