ENH: add identity kwarg to frompyfunc (#8255)

* ENH: add identity kwarg to frompyfunc

* Update umathmodule.c

* Add test, docs, and release note for identity

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

4 files changed, 48 insertions, 8 deletions

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 000000000..c0bc21b3e
--- /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 f36c6941f..cb68b8360 100644
--- a/numpy/core/_add_newdocs.py
+++ b/numpy/core/_add_newdocs.py
@@ -4180,7 +4180,7 @@ add_newdoc('numpy.core.multiarray', 'ndarray', ('view',
 
 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.
 
@@ -4195,6 +4195,13 @@ add_newdoc('numpy.core.umath', 'frompyfunc',
         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 e14006985..bad42d657 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 b7baa16e1..10a1c0803 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -2862,6 +2862,32 @@ class TestSubclass:
         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):
     """