Repository: numba/numba

Issue:
Add an option to set ufunc identity with vectorize?
Another way to write (some) custom aggregation functions, other than my suggested approach in #825, would be use plain old ufuncs.

This has the advantage of passing off the axis handling/looping to numpy's `ufunc.reduce` method. And in fact, it already _almost_ works:

```
@vectorize(['float64(float64, float64)'])
def _nanadd(a, b):
    if np.isnan(b):
        return a
    else:
        return a + b

nansum = _nanadd.reduce
```

As a bonus, this method is consistently fast. It's generally slightly (~20%) slower than bottleneck, but it has better worst case performance: `bottleneck.nansum` slows down by a factor of 4x when aggregating over a non-contiguous dimension (e.g., `axis=0` for a 2-dimensional C contiguous array), but the nansum ufunc doesn't miss a beat. I suspect this is because numpy has some clever logic for picking the optimal ordering for nested iteration for ufunc.reduce.

What's missing? There's no way to set `_nanadd.identity = 0.0`, which we would need to ensure that the calculation is correct if any of the first values are missing.

So I propose that `vectorize` should gain the optional keyword argument `identity` to set this variable. Right now it is always set to `None`: https://github.com/numba/numba/blob/1f4bee0a86edc20c9d3de07046af0d77dfa75469/numba/npyufunc/_ufunc.c#L5



Discussion:
This is a great idea.

Or perhaps make the identity attribute writable, to avoid adding another parameter to `vectorize`?

Ah, but that doesn't work due to the attribute not being writable on numpy ufuncs. Scratch that.

Ok, Numpy allows for four different identity (or unit) values:
- zero
- one
- none, meaning no unit, and the order of operations cannot be reordered (sic); for example `subtract`, `divide`, `power`
- reorderable none, no unit, and the order of operations can be reordered (sic); for example `maximum` and `minimum`

All cases except "none" allow reduction with multiple axes at once.

Interesting-- does numpy really only allow those specific values? I can see other also being potentially useful, e.g., -1 or NaN as sentinel values.

On Thu, Jan 8, 2015 at 8:54 AM, Antoine Pitrou notifications@github.com
wrote:

> Ok, Numpy allows for four different identity (or unit) values:
> - zero
> - one
> - none, meaning no unit, and the order of operations cannot be reordered (sic); for example `subtract`, `divide`, `power`
> - reorderable none, no unit, and the order of operations can be reordered (sic); for example `maximum` and `minimum`
>   All cases except "none" allow reduction with multiple axes at once.
>   ---
>   Reply to this email directly or view it on GitHub:
>   https://github.com/numba/numba/issues/829#issuecomment-69209598

I was surprised too:

https://github.com/numpy/numpy/blob/master/numpy/core/include/numpy/ufuncobject.h#L141-L142
https://github.com/numpy/numpy/blob/master/numpy/core/include/numpy/ufuncobject.h#L271-L290

I guess numba would need to define a `ReorderableNone` constant in order to allow the last case? It looks like there is no python exposure of that constant in numpy.

> I guess numba would need to define a ReorderableNone constant in order to allow the last case?

I was thinking a string (e.g. "reorderable") would be sufficient.


Resolving patch:
diff --git a/numba/npyufunc/decorators.py b/numba/npyufunc/decorators.py
--- a/numba/npyufunc/decorators.py
+++ b/numba/npyufunc/decorators.py
@@ -1,37 +1,46 @@
 from __future__ import print_function, division, absolute_import
+
+from . import _internal
 from .ufuncbuilder import UFuncBuilder, GUFuncBuilder
 
 from numba.targets.registry import TargetRegistry
 
 
-class Vectorize(object):
-    target_registry = TargetRegistry({'cpu': UFuncBuilder})
+class _BaseVectorize(object):
 
-    def __new__(cls, func, **kws):
-        target = kws.pop('target', 'cpu')
+    @classmethod
+    def get_identity(cls, kwargs):
+        return kwargs.pop('identity', None)
+
+    @classmethod
+    def get_target_implementation(cls, kwargs):
+        target = kwargs.pop('target', 'cpu')
         try:
-            imp = cls.target_registry[target]
+            return cls.target_registry[target]
         except KeyError:
             raise ValueError("Unsupported target: %s" % target)
 
-        return imp(func, kws)
+
+class Vectorize(_BaseVectorize):
+    target_registry = TargetRegistry({'cpu': UFuncBuilder})
+
+    def __new__(cls, func, **kws):
+        identity = cls.get_identity(kws)
+        imp = cls.get_target_implementation(kws)
+        return imp(func, identity, kws)
 
 
-class GUVectorize(object):
+class GUVectorize(_BaseVectorize):
     target_registry = TargetRegistry({'cpu': GUFuncBuilder})
 
     def __new__(cls, func, signature, **kws):
-        target = kws.pop('target', 'cpu')
-        try:
-            imp = cls.target_registry[target]
-        except KeyError:
-            raise ValueError("Unsupported target: %s" % target)
-
-        return imp(func, signature, kws)
+        identity = cls.get_identity(kws)
+        imp = cls.get_target_implementation(kws)
+        return imp(func, signature, identity, kws)
 
 
 def vectorize(ftylist, **kws):
-    """vectorize(ftylist[, target='cpu', [**kws]])
+    """vectorize(ftylist, target='cpu', identity=None, **kws)
 
     A decorator to create numpy ufunc object from Numba compiled code.
 
@@ -45,6 +54,11 @@ def vectorize(ftylist, **kws):
     target: str
             A string for code generation target.  Default to "cpu".
 
+    identity: int, str, or None
+        The identity (or unit) value for the element-wise function
+        being implemented.  Allowed values are None (the default), 0, 1,
+        and "reorderable".
+
     Returns
     --------
 
@@ -53,7 +67,7 @@ def vectorize(ftylist, **kws):
     Example
     -------
         @vectorize(['float32(float32, float32)',
-                    'float64(float64, float64)'])
+                    'float64(float64, float64)'], identity=1)
         def sum(a, b):
             return a + b
 
@@ -72,7 +86,7 @@ def wrap(func):
 
 
 def guvectorize(ftylist, signature, **kws):
-    """guvectorize(ftylist, signature, [, target='cpu', [**kws]])
+    """guvectorize(ftylist, signature, target='cpu', identity=None, **kws)
 
     A decorator to create numpy generialized-ufunc object from Numba compiled
     code.
@@ -84,11 +98,15 @@ def guvectorize(ftylist, signature, **kws):
         function type object or a string describing the
         function type.
 
-
     signature: str
         A NumPy generialized-ufunc signature.
         e.g. "(m, n), (n, p)->(m, p)"
 
+    identity: int, str, or None
+        The identity (or unit) value for the element-wise function
+        being implemented.  Allowed values are None (the default), 0, 1,
+        and "reorderable".
+
     target: str
             A string for code generation target.  Defaults to "cpu".
 
diff --git a/numba/npyufunc/ufuncbuilder.py b/numba/npyufunc/ufuncbuilder.py
--- a/numba/npyufunc/ufuncbuilder.py
+++ b/numba/npyufunc/ufuncbuilder.py
@@ -67,9 +67,26 @@ def compile(self, sig, locals={}, **targetoptions):
 target_registry['npyufunc'] = UFuncDispatcher
 
 
-class UFuncBuilder(object):
-    def __init__(self, py_func, targetoptions={}):
+class _BaseUFuncBuilder(object):
+
+    def parse_identity(self, identity):
+        try:
+            identity = {
+                0: _internal.PyUFunc_Zero,
+                1: _internal.PyUFunc_One,
+                None: _internal.PyUFunc_None,
+                "reorderable": _internal.PyUFunc_ReorderableNone,
+            }[identity]
+        except KeyError:
+            raise ValueError("Invalid identity value %r" % (identity,))
+        return identity
+
+
+class UFuncBuilder(_BaseUFuncBuilder):
+
+    def __init__(self, py_func, identity=None, targetoptions={}):
         self.py_func = py_func
+        self.identity = self.parse_identity(identity)
         self.nb_func = jit(target='npyufunc', **targetoptions)(py_func)
         self._sigs = []
         self._cres = {}
@@ -132,7 +149,7 @@ def build_ufunc(self):
         # there will also memory corruption. (Seems like code rewrite.)
         ufunc = _internal.fromfunc(self.py_func.__name__, self.py_func.__doc__,
                                    ptrlist, dtypelist, inct, outct, datlist,
-                                   keepalive)
+                                   keepalive, self.identity)
 
         return ufunc
 
@@ -163,11 +180,12 @@ def build(self, cres, signature):
         return dtypenums, ptr, env
 
 
-class GUFuncBuilder(object):
+class GUFuncBuilder(_BaseUFuncBuilder):
 
     # TODO handle scalar
-    def __init__(self, py_func, signature, targetoptions={}):
+    def __init__(self, py_func, signature, identity=None, targetoptions={}):
         self.py_func = py_func
+        self.identity = self.parse_identity(identity)
         self.nb_func = jit(target='npyufunc')(py_func)
         self.signature = signature
         self.sin, self.sout = parse_signature(signature)
@@ -228,7 +246,7 @@ def build_ufunc(self):
         # Pass envs to fromfuncsig to bind to the lifetime of the ufunc object
         ufunc = _internal.fromfunc(self.py_func.__name__, self.py_func.__doc__,
                                    ptrlist, dtypelist, inct, outct, datlist,
-                                   keepalive, self.signature)
+                                   keepalive, self.identity, self.signature)
         return ufunc
 
     def build(self, cres):


Test patch:
diff --git a/numba/tests/test_ufuncbuilding.py b/numba/tests/test_ufuncbuilding.py
--- a/numba/tests/test_ufuncbuilding.py
+++ b/numba/tests/test_ufuncbuilding.py
@@ -53,6 +53,7 @@ def guerror(a, b, c):
 
 
 class TestUfuncBuilding(unittest.TestCase):
+
     def test_basic_ufunc(self):
         ufb = UFuncBuilder(add)
         cres = ufb.add("int32(int32, int32)")
@@ -92,6 +93,7 @@ def test_ufunc_forceobj(self):
 
 
 class TestGUfuncBuilding(unittest.TestCase):
+
     def test_basic_gufunc(self):
         gufb = GUFuncBuilder(guadd, "(x, y),(x, y)->(x, y)")
         cres = gufb.add("void(int32[:,:], int32[:,:], int32[:,:])")
@@ -173,6 +175,21 @@ def test_vectorize_error_in_objectmode(self):
             self.assertIn("I'm here", err)
         self.assertTrue(numpy.all(b == numpy.array([1, 2, 0, 4])))
 
+    def test_vectorize_identity(self):
+        sig = 'int32(int32, int32)'
+        for identity in (0, 1, None, 'reorderable'):
+            ufunc = vectorize([sig], identity=identity)(add)
+            expected = None if identity == 'reorderable' else identity
+            self.assertEqual(ufunc.identity, expected)
+        # Default value is None
+        ufunc = vectorize([sig])(add)
+        self.assertIs(ufunc.identity, None)
+        # Invalid values
+        with self.assertRaises(ValueError):
+            vectorize([sig], identity='none')(add)
+        with self.assertRaises(ValueError):
+            vectorize([sig], identity=2)(add)
+
     def test_guvectorize(self):
         ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],
                             "(x,y),(x,y)->(x,y)")(guadd)
@@ -195,6 +212,21 @@ def test_guvectorize_error_in_objectmode(self):
         with self.assertRaises(MyException):
             ufunc(a, a)
 
+    def test_guvectorize_identity(self):
+        args = (['(int32[:,:], int32[:,:], int32[:,:])'], "(x,y),(x,y)->(x,y)")
+        for identity in (0, 1, None, 'reorderable'):
+            ufunc = guvectorize(*args, identity=identity)(guadd)
+            expected = None if identity == 'reorderable' else identity
+            self.assertEqual(ufunc.identity, expected)
+        # Default value is None
+        ufunc = guvectorize(*args)(guadd)
+        self.assertIs(ufunc.identity, None)
+        # Invalid values
+        with self.assertRaises(ValueError):
+            guvectorize(*args, identity='none')(add)
+        with self.assertRaises(ValueError):
+            guvectorize(*args, identity=2)(add)
+
 
 if __name__ == '__main__':
     unittest.main()