Add test, docs, and release note for identity

numpy · seberg · Jan 17, 2020 · Nov 8, 2016 · Dec 4, 2019 · Dec 4, 2019
commit 2e43d9d6638401b3dd8bc40a63b08602271c772c
diff --git a/doc/release/upcoming_changes/8255.new_feature.rst 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
@@ -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/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):
     """