diff --git a/doc/release/upcoming_changes/8255.new_feature.rst b/doc/release/upcoming_changes/8255.new_feature.rst new file mode 100644 index 000000000000..c0bc21b3ef1e --- /dev/null +++ b/doc/release/upcoming_changes/8255.new_feature.rst @@ -0,0 +1,5 @@ +`numpy.frompyfunc` now accepts an identity argument +--------------------------------------------------- +This allows the :attr:`numpy.ufunc.identity` attribute to be set on the +resulting ufunc, meaning it can be used for empty and multi-dimensional +calls to :meth:`numpy.ufunc.reduce`. diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py index 2f12739040bb..8993424fe494 100644 --- a/numpy/core/_add_newdocs.py +++ b/numpy/core/_add_newdocs.py @@ -4230,7 +4230,7 @@ add_newdoc('numpy.core.umath', 'frompyfunc', """ - frompyfunc(func, nin, nout) + frompyfunc(func, nin, nout, *[, identity]) Takes an arbitrary Python function and returns a NumPy ufunc. @@ -4245,6 +4245,13 @@ The number of input arguments. nout : int The number of objects returned by `func`. + identity : object, optional + The value to use for the `~numpy.ufunc.identity` attribute of the resulting + object. If specified, this is equivalent to setting the underlying + C ``identity`` field to ``PyUFunc_IdentityValue``. + If omitted, the identity is set to ``PyUFunc_None``. Note that this is + _not_ equivalent to setting the identity to ``None``, which implies the + operation is reorderable. Returns ------- diff --git a/numpy/core/src/umath/umathmodule.c b/numpy/core/src/umath/umathmodule.c index 6ec47437646d..3d41d8f813ad 100644 --- a/numpy/core/src/umath/umathmodule.c +++ b/numpy/core/src/umath/umathmodule.c @@ -70,9 +70,7 @@ object_ufunc_loop_selector(PyUFuncObject *ufunc, } PyObject * -ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *NPY_UNUSED(kwds)) { - /* Keywords are ignored for now */ - +ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *kwds) { PyObject *function, *pyname = NULL; int nin, nout, i, nargs; PyUFunc_PyFuncData *fdata; @@ -81,14 +79,18 @@ ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *NPY_UNUS Py_ssize_t fname_len = -1; void * ptr, **data; int offset[2]; + PyObject *identity = NULL; /* note: not the same semantics as Py_None */ + static char *kwlist[] = {"", "nin", "nout", "identity", NULL}; - if (!PyArg_ParseTuple(args, "Oii:frompyfunc", &function, &nin, &nout)) { + if (!PyArg_ParseTupleAndKeywords(args, kwds, "Oii|$O:frompyfunc", kwlist, + &function, &nin, &nout, &identity)) { return NULL; } if (!PyCallable_Check(function)) { PyErr_SetString(PyExc_TypeError, "function must be callable"); return NULL; } + nargs = nin + nout; pyname = PyObject_GetAttrString(function, "__name__"); @@ -146,10 +148,10 @@ ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *NPY_UNUS /* Do a better job someday */ doc = "dynamic ufunc based on a python function"; - self = (PyUFuncObject *)PyUFunc_FromFuncAndData( + self = (PyUFuncObject *)PyUFunc_FromFuncAndDataAndSignatureAndIdentity( (PyUFuncGenericFunction *)pyfunc_functions, data, - types, /* ntypes */ 1, nin, nout, PyUFunc_None, - str, doc, /* unused */ 0); + types, /* ntypes */ 1, nin, nout, identity ? PyUFunc_IdentityValue : PyUFunc_None, + str, doc, /* unused */ 0, NULL, identity); if (self == NULL) { PyArray_free(ptr); diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py index e892e81d23c6..cab98b54a0b5 100644 --- a/numpy/core/tests/test_umath.py +++ b/numpy/core/tests/test_umath.py @@ -2876,6 +2876,32 @@ def __new__(subtype, shape): a = simple((3, 4)) assert_equal(a+a, a) + +class TestFrompyfunc(object): + + def test_identity(self): + def mul(a, b): + return a * b + + # with identity=value + mul_ufunc = np.frompyfunc(mul, nin=2, nout=1, identity=1) + assert_equal(mul_ufunc.reduce([2, 3, 4]), 24) + assert_equal(mul_ufunc.reduce(np.ones((2, 2)), axis=(0, 1)), 1) + assert_equal(mul_ufunc.reduce([]), 1) + + # with identity=None (reorderable) + mul_ufunc = np.frompyfunc(mul, nin=2, nout=1, identity=None) + assert_equal(mul_ufunc.reduce([2, 3, 4]), 24) + assert_equal(mul_ufunc.reduce(np.ones((2, 2)), axis=(0, 1)), 1) + assert_raises(ValueError, lambda: mul_ufunc.reduce([])) + + # with no identity (not reorderable) + mul_ufunc = np.frompyfunc(mul, nin=2, nout=1) + assert_equal(mul_ufunc.reduce([2, 3, 4]), 24) + assert_raises(ValueError, lambda: mul_ufunc.reduce(np.ones((2, 2)), axis=(0, 1))) + assert_raises(ValueError, lambda: mul_ufunc.reduce([])) + + def _check_branch_cut(f, x0, dx, re_sign=1, im_sign=-1, sig_zero_ok=False, dtype=complex): """