This tutorial assumes you are already familiar with writing a basic ufunc, as described in
.
Background
historically only worked for ufuncs with exactly one output. This covers the .sum()/.prod()/… methods built on top of it for functions like
and
.
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
. 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
in the
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:
>>> importnumpyasnp>>> np.add.reduce([1.0,2.0,3.0])# 2 inputs, 1 outputnp.float64(6.0)>>> np.modf.reduce([1.0,2.0,3.0])# 1 input, 2 outputsTraceback (most recent call last):...ValueError: reduce only supported for binary functions>>> np.divmod.reduce([1.0,2.0,3.0])# 2 inputs, 2 outputsTraceback (most recent call last):...TypeError: divmod.reduce is not supported: the resolved loop does not register a reduction loopThe first two behaviours are unchanged by this feature.
takes a single input, so it can never be reduced. Only the third case is new:
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 (
), not the legacy PyUFunc_FromFuncAndData loop table used in the previous tutorial. The ufunc itself is still created with
, 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 /*BEGINmaincomputation*/ and /*ENDmaincomputation*/.
#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. */staticintdouble_minimummaximum_loop(PyArrayMethod_Context*NPY_UNUSED(context),char*constdata[],npy_intpconstdimensions[],npy_intpconststrides[],NpyAuxData*NPY_UNUSED(auxdata)){npy_intpn=dimensions[0];char*in1=data[0],*in2=data[1];char*out1=data[2],*out2=data[3];npy_intpis1=strides[0],is2=strides[1];npy_intpos1=strides[2],os2=strides[3];for(npy_intpi=0;i<n;i++){/* BEGIN main computation */doublea=*(double*)in1;doubleb=*(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;}return0;}/* * 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. */staticintdouble_minimummaximum_reduce_loop(PyArrayMethod_Context*NPY_UNUSED(context),char*constdata[],npy_intpconstdimensions[],npy_intpconststrides[],NpyAuxData*NPY_UNUSED(auxdata)){npy_intpn=dimensions[0];char*acc_min=data[0],*acc_max=data[1],*x=data[2];char*out_min=data[3],*out_max=data[4];npy_intps_amin=strides[0],s_amax=strides[1],s_x=strides[2];npy_intps_omin=strides[3],s_omax=strides[4];for(npy_intpi=0;i<n;i++){/* BEGIN main computation */doublecur_min=*(double*)acc_min;doublecur_max=*(double*)acc_max;doubleval=*(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;}return0;}/* * 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. */staticintminimummaximum_get_reduction_loop(PyArrayMethod_Context*NPY_UNUSED(context),intNPY_UNUSED(aligned),intNPY_UNUSED(move_references),constnpy_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;return0;}/* * 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. */staticintminimummaximum_promoter(PyObject*NPY_UNUSED(ufunc),PyArray_DTypeMeta*constNPY_UNUSED(op_dtypes[]),PyArray_DTypeMeta*constsignature[],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(inti=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);return0;}staticintregister_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;}intres=PyUFunc_AddPromoter(minimummaximum,none_tuple,promoter);Py_DECREF(none_tuple);Py_DECREF(promoter);returnres;}staticPyMethodDefMinMaxMethods[]={{NULL,NULL,0,NULL}};staticstructPyModuleDefmoduledef={PyModuleDef_HEAD_INIT,"minmax",NULL,-1,MinMaxMethods,NULL,NULL,NULL,NULL};PyMODINIT_FUNCPyInit_minmax(void){import_array();import_umath();PyObject*m=PyModule_Create(&moduledef);if(m==NULL){returnNULL;}/* 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);returnNULL;}PyArray_Descr*double_descr=PyArray_DescrFromType(NPY_DOUBLE);PyArray_DTypeMeta*dt=NPY_DTYPE(double_descr);PyArray_DTypeMeta*dtypes[4]={dt,dt,dt,dt};PyType_Slotslots[]={{NPY_METH_strided_loop,(void*)&double_minimummaximum_loop},{NPY_METH_get_reduction_loop,(void*)&minimummaximum_get_reduction_loop},{0,NULL},};PyArrayMethod_Specspec={.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,};intres=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);returnNULL;}returnm;}
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"fromsetuptoolsimportsetup,Extensionfromnumpyimportget_includeminmax=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:
>>> importnumpyasnp>>> importminmax>>> 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
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
/
do.
It registers no reduction identity. Reducing an empty array therefore raises a
, 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:
staticintminimummaximum_get_multi_reduction_initials(PyArrayMethod_Context*NPY_UNUSED(context),npy_boolNPY_UNUSED(reduction_is_empty),void**initials){*(double*)initials[0]=NPY_INFINITY;/*minidentity*/*(double*)initials[1]=-NPY_INFINITY;/*maxidentity*/return1;}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}.
,
, and
are not supported for multi-output ufuncs yet, only
is.
See also
for the full
reference documentation.
Universal functions (ufunc) basics
for the Python-level behavior of
on multi-output ufuncs.