Adding a reduction loop to a ufunc

This tutorial assumes you are already familiar with writing a basic ufunc, as described in Writing your own ufunc.

Background

reduce historically only worked for ufuncs with exactly one output. This covers the .sum()/.prod()/… methods built on top of it for functions like numpy.add and numpy.multiply.

What makes those work is that the ordinary two-input/one-output elementwise loop can double as the reduction loop. NumPy points the loop’s output and its first input at the same piece of memory. It also gives that pointer a stride of zero, so the pointer never advances. The loop then reads the running result and writes the updated result back to that same location as it walks over the array.

That trick does not generalize to ufuncs with more than one output, such as numpy.divmod. Their forward loop has a two-in/nout-out signature. There is no natural way to pair each output up with an input that could hold its running result.

To support reduce for such ufuncs, an ArrayMethod can register a dedicated reduction loop through the NPY_METH_get_reduction_loop slot. That loop has an (nout + 1)-in/nout-out signature. It takes the current per-output accumulators plus one streamed input element, and writes the updated accumulators. We use this signature because it collapses down to the usual two-input/one-output case when nout is 1. See NPY_METH_get_reduction_loop in the ArrayMethod API reference for the exact data/stride layout.

Only the restriction on the number of outputs is lifted. The number of inputs is unchanged: reduce still requires exactly two, as do accumulate and reduceat. Reducing repeatedly combines an accumulated result with one new element, which only makes sense for a ufunc taking two inputs.

The three cases look like this:

>>> import numpy as np
>>> np.add.reduce([1.0, 2.0, 3.0])  # 2 inputs, 1 output
np.float64(6.0)
>>> np.modf.reduce([1.0, 2.0, 3.0])  # 1 input, 2 outputs
Traceback (most recent call last):
    ...
ValueError: reduce only supported for binary functions
>>> np.divmod.reduce([1.0, 2.0, 3.0])  # 2 inputs, 2 outputs
Traceback (most recent call last):
    ...
TypeError: divmod.reduce is not supported: the resolved loop does not register a reduction loop

The first two behaviours are unchanged by this feature. numpy.modf takes a single input, so it can never be reduced. Only the third case is new: numpy.divmod has the right number of inputs, and fails purely because no reduction loop is registered for it. The rest of this tutorial shows how to write and register such a loop.

Example: a two-output minimummaximum ufunc

The example below defines a minimummaximum ufunc for the 'f8' (double) dtype only. It computes the running minimum and maximum in a single pass, so minimummaximum.reduce(a) returns the pair (a.min(), a.max()) using one loop over a rather than two.

NPY_METH_get_reduction_loop is an ArrayMethod slot. The ufunc’s loop must therefore be registered through the PyArrayMethod_Spec API (PyUFunc_AddLoopFromSpec), not the legacy PyUFunc_FromFuncAndData loop table used in the previous tutorial. The ufunc itself is still created with PyUFunc_FromFuncAndData, just without any loops attached yet. The loop is added separately, via its ArrayMethod spec.

The forward (elementwise) loop is 2-in/2-out:

(a, b) -> (min(a, b), max(a, b))

and the reduction loop is 3-in/2-out (nout = 2 outputs, so nout + 1 = 3 inputs):

(acc_min, acc_max, x) -> (min(acc_min, x), max(acc_max, x))

The place in the code corresponding to the actual computations is marked with /* BEGIN main computation */ and /* END main computation */.

#define PY_SSIZE_T_CLEAN
#define NPY_TARGET_VERSION NPY_2_6_API_VERSION
#define NPY_NO_DEPRECATED_API NPY_API_VERSION
#include <Python.h>
#include "numpy/ndarraytypes.h"
#include "numpy/ufuncobject.h"
#include "numpy/dtype_api.h"

/*
 * minmax.c
 * A minimal ufunc for a single dtype ('f8') that computes the
 * running minimum and maximum in one pass, and registers a
 * reduction loop for it so that `minimummaximum.reduce(a)` computes
 * the actual minmax.
 */

static int
double_minimummaximum_loop(PyArrayMethod_Context *NPY_UNUSED(context),
        char *const data[], npy_intp const dimensions[],
        npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata))
{
    npy_intp n = dimensions[0];
    char *in1 = data[0], *in2 = data[1];
    char *out1 = data[2], *out2 = data[3];
    npy_intp is1 = strides[0], is2 = strides[1];
    npy_intp os1 = strides[2], os2 = strides[3];

    for (npy_intp i = 0; i < n; i++) {
        /* BEGIN main computation */
        double a = *(double *)in1;
        double b = *(double *)in2;
        *(double *)out1 = a < b ? a : b;
        *(double *)out2 = a < b ? b : a;
        /* END main computation */
        in1 += is1; in2 += is2; out1 += os1; out2 += os2;
    }
    return 0;
}

/*
 * The reduce machinery points each out_i at the same memory as the
 * matching acc_i (with stride 0), so this both reads the running
 * accumulators and writes the new ones.
 */
static int
double_minimummaximum_reduce_loop(PyArrayMethod_Context *NPY_UNUSED(context),
        char *const data[], npy_intp const dimensions[],
        npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata))
{
    npy_intp n = dimensions[0];
    char *acc_min = data[0], *acc_max = data[1], *x = data[2];
    char *out_min = data[3], *out_max = data[4];
    npy_intp s_amin = strides[0], s_amax = strides[1], s_x = strides[2];
    npy_intp s_omin = strides[3], s_omax = strides[4];

    for (npy_intp i = 0; i < n; i++) {
        /* BEGIN main computation */
        double cur_min = *(double *)acc_min;
        double cur_max = *(double *)acc_max;
        double val = *(double *)x;
        *(double *)out_min = val < cur_min ? val : cur_min;
        *(double *)out_max = val > cur_max ? val : cur_max;
        /* END main computation */
        acc_min += s_amin; acc_max += s_amax; x += s_x;
        out_min += s_omin; out_max += s_omax;
    }
    return 0;
}

/*
 * get_reduction_loop is called by the reduce machinery in place of
 * get_strided_loop.  It hands back the (nout+1)->nout loop above
 * instead of the forward elementwise one.
 */
static int
minimummaximum_get_reduction_loop(PyArrayMethod_Context *NPY_UNUSED(context),
        int NPY_UNUSED(aligned), int NPY_UNUSED(move_references),
        const npy_intp *NPY_UNUSED(strides),
        PyArrayMethod_StridedLoop **out_loop,
        NpyAuxData **out_transferdata,
        NPY_ARRAYMETHOD_FLAGS *flags)
{
    *out_loop = &double_minimummaximum_reduce_loop;
    *out_transferdata = NULL;
    *flags = NPY_METH_NO_FLOATINGPOINT_ERRORS;
    return 0;
}

/*
 * Only one loop is registered ('f8'), so the promoter always
 * requests float64 for every input/output that isn't already
 * pinned by `signature` (e.g. via `dtype=`). This is what lets
 * plain Python floats/ints, or other numeric dtypes, be used
 * directly. Without it, only exact `np.float64` inputs would
 * match the loop registered below. A ufunc that registers loops
 * for several dtypes would instead pick the common dtype of its
 * inputs here (see `PyArray_PromoteDTypeSequence`). This one only
 * ever has a single dtype to offer.
 */
static int
minimummaximum_promoter(PyObject *NPY_UNUSED(ufunc),
        PyArray_DTypeMeta *const NPY_UNUSED(op_dtypes[]),
        PyArray_DTypeMeta *const signature[],
        PyArray_DTypeMeta *new_op_dtypes[])
{
    PyArray_Descr *double_descr = PyArray_DescrFromType(NPY_DOUBLE);
    if (double_descr == NULL) {
        return -1;
    }
    PyArray_DTypeMeta *double_dt = NPY_DTYPE(double_descr);

    for (int i = 0; i < 4; i++) {
        PyArray_DTypeMeta *dt = signature[i] != NULL ? signature[i] : double_dt;
        Py_INCREF(dt);
        new_op_dtypes[i] = dt;
    }
    Py_DECREF(double_descr);
    return 0;
}

static int
register_minimummaximum_promoter(PyObject *minimummaximum)
{
    /* All-None tuple -> catch-all fallback (a concrete loop always wins). */
    PyObject *none_tuple = PyTuple_Pack(
            4, Py_None, Py_None, Py_None, Py_None);
    if (none_tuple == NULL) {
        return -1;
    }
    PyObject *promoter = PyCapsule_New(
            (void *)&minimummaximum_promoter,
            "numpy._ufunc_promoter", NULL);
    if (promoter == NULL) {
        Py_DECREF(none_tuple);
        return -1;
    }
    int res = PyUFunc_AddPromoter(minimummaximum, none_tuple, promoter);
    Py_DECREF(none_tuple);
    Py_DECREF(promoter);
    return res;
}

static PyMethodDef MinMaxMethods[] = {
    {NULL, NULL, 0, NULL}
};

static struct PyModuleDef moduledef = {
    PyModuleDef_HEAD_INIT, "minmax", NULL, -1, MinMaxMethods,
    NULL, NULL, NULL, NULL
};

PyMODINIT_FUNC PyInit_minmax(void)
{
    import_array();
    import_umath();

    PyObject *m = PyModule_Create(&moduledef);
    if (m == NULL) {
        return NULL;
    }

    /* A 2-in/2-out ufunc. Its loops are registered separately below. */
    PyObject *minimummaximum = PyUFunc_FromFuncAndData(
            NULL, NULL, NULL, 0, 2, 2, PyUFunc_None,
            "minimummaximum", "minimummaximum_docstring", 0);
    if (minimummaximum == NULL) {
        Py_DECREF(m);
        return NULL;
    }

    PyArray_Descr *double_descr = PyArray_DescrFromType(NPY_DOUBLE);
    PyArray_DTypeMeta *dt = NPY_DTYPE(double_descr);
    PyArray_DTypeMeta *dtypes[4] = {dt, dt, dt, dt};

    PyType_Slot slots[] = {
        {NPY_METH_strided_loop, (void *)&double_minimummaximum_loop},
        {NPY_METH_get_reduction_loop,
         (void *)&minimummaximum_get_reduction_loop},
        {0, NULL},
    };

    PyArrayMethod_Spec spec = {
        .name = "double_minimummaximum",
        .nin = 2,
        .nout = 2,
        .casting = NPY_NO_CASTING,
        /* Needed for axis=None / multi-axis reductions. */
        .flags = NPY_METH_IS_REORDERABLE,
        .dtypes = dtypes,
        .slots = slots,
    };

    int res = PyUFunc_AddLoopFromSpec(minimummaximum, &spec);
    Py_DECREF(double_descr);
    if (res < 0 || register_minimummaximum_promoter(minimummaximum) < 0
            || PyModule_AddObject(m, "minimummaximum", minimummaximum) < 0) {
        Py_XDECREF(minimummaximum);
        Py_DECREF(m);
        return NULL;
    }
    return m;
}

As with the ufuncs in the previous tutorial, the module needs to declare a dependency on NumPy to build:

meson

Sample pyproject.toml and meson.build.

[project]
name = "minmax"
dependencies = ["numpy"]
version = "0.1"

[build-system]
requires = ["meson-python", "numpy"]
build-backend = "mesonpy"

project('minmax', 'c')

py = import('python').find_installation()
np_dep = dependency('numpy')

sources = files('minmax.c')

extension_module = py.extension_module(
  'minmax',
  sources,
  dependencies: [np_dep],
  install: true,
)

setuptools

Sample pyproject.toml and setup.py.

[project]
name = "minmax"
dependencies = ["numpy"]
version = "0.1"

[build-system]
requires = ["setuptools", "numpy"]
build-backend = "setuptools.build_meta"

from setuptools import setup, Extension
from numpy import get_include

minmax = Extension('minmax',
                   sources=['minmax.c'],
                   include_dirs=[get_include()])

setup(name='minmax', version='1.0', ext_modules=[minmax])

After the above has been installed, it can be imported and used as follows. Plain Python floats and ints work directly, thanks to the promoter. No explicit np.float64(...) wrapping is needed:

>>> import numpy as np
>>> import minmax
>>> minmax.minimummaximum(3.0, 5.0)
(np.float64(3.0), np.float64(5.0))
>>> minmax.minimummaximum(3, 5)
(np.float64(3.0), np.float64(5.0))
>>> a = [3.0, 1.0, 4.0, 1.0, 5.0, 9.0, 2.0, 6.0]
>>> minmax.minimummaximum.reduce(a)
(np.float64(1.0), np.float64(9.0))
>>> b = np.array(a).reshape(2, 4)
>>> minmax.minimummaximum.reduce(b, axis=None)
(np.float64(1.0), np.float64(9.0))
>>> minmax.minimummaximum.reduce(b, axis=1)
(array([1., 2.]), array([4., 9.]))

Note that minimummaximum.reduce returns a tuple with one array per output, rather than a single array, so it is usually unpacked as lo, hi = minimummaximum.reduce(a). The initial argument of a multi-output reduce accepts either a single scalar, which seeds every output, or a tuple with one value per output. The entries of such a tuple cannot be None; a plain initial=None still unsets the initial value for the whole reduction.

Limitations

This example is intentionally minimal and leaves out several things a real-world implementation, like the built-in numpy.add or a fuller minimummaximum, would normally have:

• The promoter always resolves to 'f8', regardless of the input dtypes. Integer input arrays are therefore silently cast to float64 instead of keeping their own dtype. A multi-dtype ufunc would register a loop per dtype and pick the common one, as the built-in numpy.minimum/numpy.maximum do.

• It registers no reduction identity. Reducing an empty array therefore raises a ValueError, the same as numpy.maximum.reduce. A multi-output ufunc can provide per-output identities with the NPY_METH_get_multi_reduction_initials slot, the multi-output version of NPY_METH_get_reduction_initial. Its function fills one initial value per output, for example +inf for the running minimum and -inf for the running maximum. Reducing an empty array, or reducing with a where= mask, then returns those identities instead of raising:

  static int
  minimummaximum_get_multi_reduction_initials(
          PyArrayMethod_Context *NPY_UNUSED(context),
          npy_bool NPY_UNUSED(reduction_is_empty), void **initials)
  {
      *(double *)initials[0] = NPY_INFINITY;   /* min identity */
      *(double *)initials[1] = -NPY_INFINITY;  /* max identity */
      return 1;
  }
  

  It is registered alongside the reduction loop, by adding this entry to the slots array above: {NPY_METH_get_multi_reduction_initials, (void *)&minimummaximum_get_multi_reduction_initials}.

• accumulate, reduceat, and at are not supported for multi-output ufuncs yet, only reduce is.

See also

• ArrayMethod API reference for the full NPY_METH_get_reduction_loop reference documentation.

• Universal functions (ufunc) basics for the Python-level behavior of reduce on multi-output ufuncs.