Merge pull request #8955 from eric-wieser/obj-identity

ENH: Allow ufunc.identity to be any python object

7 files changed, 127 insertions, 20 deletions

diff --git a/doc/release/1.16.0-notes.rst b/doc/release/1.16.0-notes.rst
index 3e00e4410..7f2a843eb 100644
--- a/doc/release/1.16.0-notes.rst
+++ b/doc/release/1.16.0-notes.rst
@@ -109,8 +109,13 @@ for unraveling. ``dims`` remains supported, but is now deprecated.
 C API changes
 =============
 
-The :c:data:`NPY_API_VERSION` was incremented to 0x0000D since
-``core_dim_flags`` and ``core_dim_sizes`` were added to :c:type:`PyUFuncObject`.
+The :c:data:`NPY_API_VERSION` was incremented to 0x0000D, due to the addition
+of:
+
+* :c:member:`PyUFuncObject.core_dim_flags`
+* :c:member:`PyUFuncObject.core_dim_sizes`
+* :c:member:`PyUFuncObject.identity_value`
+* :c:function:`PyUFunc_FromFuncAndDataAndSignatureAndIdentity`
 
 New Features
 ============
@@ -287,6 +292,25 @@ and other path-like objects in addition to a file object. Furthermore, the
 ``np.load`` function now also supports path-like objects when
 using memory mapping (``mmap_mode`` keyword argument).
 
+Better behaviour of ufunc identities during reductions
+------------------------------------------------------
+Universal functions have an ``.identity`` which is used when ``.reduce`` is
+called on an empty axis.
+
+As of this release, the logical binary ufuncs, `logical_and`, `logical_or`,
+and `logical_xor`, now have ``identity``s of type `bool`, where previously they
+were of type `int`. This restores the 1.14 behavior of getting ``bool``s when
+reducing empty object arrays with these ufuncs, while also keeping the 1.15
+behavior of getting ``int``s when reducing empty object arrays with arithmetic
+ufuncs like ``add`` and ``multiply``.
+
+Additionally, `logaddexp` now has an identity of ``-inf``, allowing it to be
+called on empty sequences, where previously it could not be.
+
+This is possible thanks to the new
+:c:function:`PyUFunc_FromFuncAndDataAndSignatureAndIdentity`, which allows
+arbitrary values to be used as identities now.
+
 
 Changes
 =======
diff --git a/doc/source/reference/c-api.ufunc.rst b/doc/source/reference/c-api.ufunc.rst
index 07c7b0c80..0499ccf5b 100644
--- a/doc/source/reference/c-api.ufunc.rst
+++ b/doc/source/reference/c-api.ufunc.rst
@@ -169,8 +169,12 @@ Functions
     :param identity:
 
         Either :c:data:`PyUFunc_One`, :c:data:`PyUFunc_Zero`,
-        :c:data:`PyUFunc_None`. This specifies what should be returned when
+        :c:data:`PyUFunc_MinusOne`, or :c:data:`PyUFunc_None`.
+        This specifies what should be returned when
         an empty array is passed to the reduce method of the ufunc.
+        The special value :c:data:`PyUFunc_IdentityValue` may only be used with
+        the :c:func:`PyUFunc_FromFuncAndDataAndSignatureAndIdentity` method, to
+        allow an arbitrary python object to be used as the identity.
 
     :param name:
         The name for the ufunc as a ``NULL`` terminated string.  Specifying
@@ -206,6 +210,21 @@ Functions
         to calling PyUFunc_FromFuncAndData. A copy of the string is made,
         so the passed in buffer can be freed.
 
+.. c:function:: PyObject* PyUFunc_FromFuncAndDataAndSignatureAndIdentity(
+        PyUFuncGenericFunction *func, void **data, char *types, int ntypes, \
+        int nin, int nout, int identity, char *name, char *doc, int unused, char *signature,
+        PyObject *identity_value)
+
+   This function is very similar to `PyUFunc_FromFuncAndDataAndSignature` above,
+   but has an extra *identity_value* argument, to define an arbitrary identity
+   for the ufunc when ``identity`` is passed as ``PyUFunc_IdentityValue``.
+
+   :param identity_value:
+        The identity for the new gufunc. Must be passed as ``NULL`` unless the
+        ``identity`` argument is ``PyUFunc_IdentityValue``. Setting it to NULL
+        is equivalent to calling PyUFunc_FromFuncAndDataAndSignature.
+
+
 .. c:function:: int PyUFunc_RegisterLoopForType( \
         PyUFuncObject* ufunc, int usertype, PyUFuncGenericFunction function, \
         int* arg_types, void* data)
diff --git a/numpy/core/code_generators/generate_umath.py b/numpy/core/code_generators/generate_umath.py
index 199ad831b..9d4e72c0e 100644
--- a/numpy/core/code_generators/generate_umath.py
+++ b/numpy/core/code_generators/generate_umath.py
@@ -10,11 +10,14 @@ sys.path.insert(0, os.path.dirname(__file__))
 import ufunc_docstrings as docstrings
 sys.path.pop(0)
 
-Zero = "PyUFunc_Zero"
-One = "PyUFunc_One"
-None_ = "PyUFunc_None"
-AllOnes = "PyUFunc_MinusOne"
-ReorderableNone = "PyUFunc_ReorderableNone"
+Zero = "PyInt_FromLong(0)"
+One = "PyInt_FromLong(1)"
+True_ = "(Py_INCREF(Py_True), Py_True)"
+False_ = "(Py_INCREF(Py_False), Py_False)"
+None_ = object()
+AllOnes = "PyInt_FromLong(-1)"
+MinusInfinity = 'PyFloat_FromDouble(-NPY_INFINITY)'
+ReorderableNone = "(Py_INCREF(Py_None), Py_None)"
 
 # Sentinel value to specify using the full type description in the
 # function name
@@ -458,7 +461,7 @@ defdict = {
           [TypeDescription('O', FullTypeDescr, 'OO', 'O')],
           ),
 'logical_and':
-    Ufunc(2, 1, One,
+    Ufunc(2, 1, True_,
           docstrings.get('numpy.core.umath.logical_and'),
           'PyUFunc_SimpleBinaryComparisonTypeResolver',
           TD(nodatetime_or_obj, out='?', simd=[('avx2', ints)]),
@@ -472,14 +475,14 @@ defdict = {
           TD(O, f='npy_ObjectLogicalNot'),
           ),
 'logical_or':
-    Ufunc(2, 1, Zero,
+    Ufunc(2, 1, False_,
           docstrings.get('numpy.core.umath.logical_or'),
           'PyUFunc_SimpleBinaryComparisonTypeResolver',
           TD(nodatetime_or_obj, out='?', simd=[('avx2', ints)]),
           TD(O, f='npy_ObjectLogicalOr'),
           ),
 'logical_xor':
-    Ufunc(2, 1, Zero,
+    Ufunc(2, 1, False_,
           docstrings.get('numpy.core.umath.logical_xor'),
           'PyUFunc_SimpleBinaryComparisonTypeResolver',
           TD(nodatetime_or_obj, out='?'),
@@ -514,7 +517,7 @@ defdict = {
           TD(O, f='npy_ObjectMin')
           ),
 'logaddexp':
-    Ufunc(2, 1, None,
+    Ufunc(2, 1, MinusInfinity,
           docstrings.get('numpy.core.umath.logaddexp'),
           None,
           TD(flts, f="logaddexp", astype={'e':'f'})
@@ -1048,18 +1051,38 @@ def make_ufuncs(funcdict):
         # do not play well with \n
         docstring = '\\n\"\"'.join(docstring.split(r"\n"))
         fmt = textwrap.dedent("""\
-            f = PyUFunc_FromFuncAndData(
+            identity = {identity_expr};
+            if ({has_identity} && identity == NULL) {{
+                return -1;
+            }}
+            f = PyUFunc_FromFuncAndDataAndSignatureAndIdentity(
                 {name}_functions, {name}_data, {name}_signatures, {nloops},
                 {nin}, {nout}, {identity}, "{name}",
-                "{doc}", 0
+                "{doc}", 0, NULL, identity
             );
+            if ({has_identity}) {{
+                Py_DECREF(identity);
+            }}
             if (f == NULL) {{
                 return -1;
-            }}""")
-        mlist.append(fmt.format(
+            }}
+        """)
+        args = dict(
             name=name, nloops=len(uf.type_descriptions),
-            nin=uf.nin, nout=uf.nout, identity=uf.identity, doc=docstring
-        ))
+            nin=uf.nin, nout=uf.nout,
+            has_identity='0' if uf.identity is None_ else '1',
+            identity='PyUFunc_IdentityValue',
+            identity_expr=uf.identity,
+            doc=docstring
+        )
+
+        # Only PyUFunc_None means don't reorder - we pass this using the old
+        # argument
+        if uf.identity is None_:
+            args['identity'] = 'PyUFunc_None'
+            args['identity_expr'] = 'NULL'
+
+        mlist.append(fmt.format(**args))
         if uf.typereso is not None:
             mlist.append(
                 r"((PyUFuncObject *)f)->type_resolver = &%s;" % uf.typereso)
@@ -1087,7 +1110,7 @@ def make_code(funcdict, filename):
 
     static int
     InitOperators(PyObject *dictionary) {
-        PyObject *f;
+        PyObject *f, *identity;
 
     %s
     %s
diff --git a/numpy/core/code_generators/numpy_api.py b/numpy/core/code_generators/numpy_api.py
index d8a9ee6b4..fdf97ac00 100644
--- a/numpy/core/code_generators/numpy_api.py
+++ b/numpy/core/code_generators/numpy_api.py
@@ -402,6 +402,9 @@ ufunc_funcs_api = {
     # End 1.7 API
     'PyUFunc_RegisterLoopForDescr':             (41,),
     # End 1.8 API
+    'PyUFunc_FromFuncAndDataAndSignatureAndIdentity':
+                                                (42,),
+    # End 1.16 API
 }
 
 # List of all the dicts which define the C API
diff --git a/numpy/core/include/numpy/ufuncobject.h b/numpy/core/include/numpy/ufuncobject.h
index 85f8a6c08..90d837a9b 100644
--- a/numpy/core/include/numpy/ufuncobject.h
+++ b/numpy/core/include/numpy/ufuncobject.h
@@ -223,7 +223,8 @@ typedef struct _tagPyUFuncObject {
          */
         npy_uint32 *core_dim_flags;
 
-
+        /* Identity for reduction, when identity == PyUFunc_IdentityValue */
+        PyObject *identity_value;
 
 } PyUFuncObject;
 
@@ -299,6 +300,12 @@ typedef struct _tagPyUFuncObject {
  * This case allows reduction with multiple axes at once.
  */
 #define PyUFunc_ReorderableNone -2
+/*
+ * UFunc unit is in identity_value, and the order of operations can be reordered
+ * This case allows reduction with multiple axes at once.
+ */
+#define PyUFunc_IdentityValue -3
+
 
 #define UFUNC_REDUCE 0
 #define UFUNC_ACCUMULATE 1
diff --git a/numpy/core/src/umath/ufunc_object.c b/numpy/core/src/umath/ufunc_object.c
index e60c734ec..1fe8745a0 100644
--- a/numpy/core/src/umath/ufunc_object.c
+++ b/numpy/core/src/umath/ufunc_object.c
@@ -2453,6 +2453,11 @@ _get_identity(PyUFuncObject *ufunc, npy_bool *reorderable) {
         *reorderable = 0;
         Py_RETURN_NONE;
 
+    case PyUFunc_IdentityValue:
+        *reorderable = 1;
+        Py_INCREF(ufunc->identity_value);
+        return ufunc->identity_value;
+
     default:
         PyErr_Format(PyExc_ValueError,
                 "ufunc %s has an invalid identity", ufunc_get_name_cstr(ufunc));
@@ -4833,6 +4838,20 @@ PyUFunc_FromFuncAndDataAndSignature(PyUFuncGenericFunction *func, void **data,
                                      const char *name, const char *doc,
                                      int unused, const char *signature)
 {
+    return PyUFunc_FromFuncAndDataAndSignatureAndIdentity(
+        func, data, types, ntypes, nin, nout, identity, name, doc,
+        unused, signature, NULL);
+}
+
+/*UFUNC_API*/
+NPY_NO_EXPORT PyObject *
+PyUFunc_FromFuncAndDataAndSignatureAndIdentity(PyUFuncGenericFunction *func, void **data,
+                                     char *types, int ntypes,
+                                     int nin, int nout, int identity,
+                                     const char *name, const char *doc,
+                                     int unused, const char *signature,
+                                     PyObject *identity_value)
+{
     PyUFuncObject *ufunc;
     if (nin + nout > NPY_MAXARGS) {
         PyErr_Format(PyExc_ValueError,
@@ -4853,6 +4872,10 @@ PyUFunc_FromFuncAndDataAndSignature(PyUFuncGenericFunction *func, void **data,
     ufunc->nout = nout;
     ufunc->nargs = nin+nout;
     ufunc->identity = identity;
+    if (ufunc->identity == PyUFunc_IdentityValue) {
+        Py_INCREF(identity_value);
+    }
+    ufunc->identity_value = identity_value;
 
     ufunc->functions = func;
     ufunc->data = data;
@@ -4874,6 +4897,7 @@ PyUFunc_FromFuncAndDataAndSignature(PyUFuncGenericFunction *func, void **data,
 
     ufunc->op_flags = PyArray_malloc(sizeof(npy_uint32)*ufunc->nargs);
     if (ufunc->op_flags == NULL) {
+        Py_DECREF(ufunc);
         return PyErr_NoMemory();
     }
     memset(ufunc->op_flags, 0, sizeof(npy_uint32)*ufunc->nargs);
@@ -5230,6 +5254,9 @@ ufunc_dealloc(PyUFuncObject *ufunc)
     PyArray_free(ufunc->op_flags);
     Py_XDECREF(ufunc->userloops);
     Py_XDECREF(ufunc->obj);
+    if (ufunc->identity == PyUFunc_IdentityValue) {
+        Py_DECREF(ufunc->identity_value);
+    }
     PyArray_free(ufunc);
 }
 
diff --git a/numpy/core/tests/test_umath.py b/numpy/core/tests/test_umath.py
index bd7985dfb..f32c2968f 100644
--- a/numpy/core/tests/test_umath.py
+++ b/numpy/core/tests/test_umath.py
@@ -685,6 +685,10 @@ class TestLogAddExp(_FilterInvalids):
         assert_(np.isnan(np.logaddexp(0, np.nan)))
         assert_(np.isnan(np.logaddexp(np.nan, np.nan)))
 
+    def test_reduce(self):
+        assert_equal(np.logaddexp.identity, -np.inf)
+        assert_equal(np.logaddexp.reduce([]), -np.inf)
+
 
 class TestLog1p(object):
     def test_log1p(self):