Repository: numba/numba

Issue:
Support for numpy.round/round_/around
In Numpy, `round`/`round_`/`around` are all aliases to the same ufunc, which does element-wise rounding in the same fashion as the Python standard `round()` function:

http://docs.scipy.org/doc/numpy/reference/generated/numpy.round_.html

Now that we have support for `round()` in Numba, it would be nice to support the Numpy-equivalent as well.



Discussion:
Interesting. `round()` is actually not a ufunc but redirects to the array object's `.round()` method.

Actually, I recall adding plenty of support for ufuncs, and a way to find out which ufuncs (there is a list somewhere in the docs) was looking for symbols in the numpy namespace that were ufunc objects. So it is very likely that it is not implemented because of that. I wonder how many of those "ufunc-like" functions  are missing that would be quite easy to support.

Also of note, some of the support for ufuncs _rely_ on the ufunc objects themselves, as their type inference is based on some information ufuncs have with them, but adding support for other "ufunc-like" functions wouldn't be difficult if that data is defined somewhere else.


Resolving patch:
diff --git a/numba/targets/arrayobj.py b/numba/targets/arrayobj.py
--- a/numba/targets/arrayobj.py
+++ b/numba/targets/arrayobj.py
@@ -6,6 +6,7 @@
 from __future__ import print_function, absolute_import, division
 
 from functools import reduce
+import math
 
 import llvmlite.llvmpy.core as lc
 
@@ -622,6 +623,93 @@ def array_argmax_impl(arry):
     return context.compile_internal(builder, array_argmax_impl, sig, args)
 
 
+def _np_round_intrinsic(tp):
+    # np.round() always rounds half to even
+    return "llvm.rint.f%d" % (tp.bitwidth,)
+
+def _np_round_float(context, builder, tp, val):
+    llty = context.get_value_type(tp)
+    module = cgutils.get_module(builder)
+    fnty = lc.Type.function(llty, [llty])
+    fn = module.get_or_insert_function(fnty, name=_np_round_intrinsic(tp))
+    return builder.call(fn, (val,))
+
+@builtin
+@implement(numpy.round, types.Kind(types.Float))
+def scalar_round_unary(context, builder, sig, args):
+    return _np_round_float(context, builder, sig.args[0], args[0])
+
+@builtin
+@implement(numpy.round, types.Kind(types.Integer))
+def scalar_round_unary(context, builder, sig, args):
+    return args[0]
+
+@builtin
+@implement(numpy.round, types.Kind(types.Complex))
+def scalar_round_unary_complex(context, builder, sig, args):
+    fltty = sig.args[0].underlying_float
+    cplx_cls = context.make_complex(sig.args[0])
+    z = cplx_cls(context, builder, args[0])
+    z.real = _np_round_float(context, builder, fltty, z.real)
+    z.imag = _np_round_float(context, builder, fltty, z.imag)
+    return z._getvalue()
+
+@builtin
+@implement(numpy.round, types.Kind(types.Float), types.Kind(types.Integer))
+@implement(numpy.round, types.Kind(types.Integer), types.Kind(types.Integer))
+def scalar_round_binary_float(context, builder, sig, args):
+    def round_ndigits(x, ndigits):
+        if math.isinf(x) or math.isnan(x):
+            return x
+
+        # NOTE: this is CPython's algorithm, but perhaps this is overkill
+        # when emulating Numpy's behaviour.
+        if ndigits >= 0:
+            if ndigits > 22:
+                # pow1 and pow2 are each safe from overflow, but
+                # pow1*pow2 ~= pow(10.0, ndigits) might overflow.
+                pow1 = 10.0 ** (ndigits - 22)
+                pow2 = 1e22
+            else:
+                pow1 = 10.0 ** ndigits
+                pow2 = 1.0
+            y = (x * pow1) * pow2
+            if math.isinf(y):
+                return x
+            return (numpy.round(y) / pow2) / pow1
+
+        else:
+            pow1 = 10.0 ** (-ndigits)
+            y = x / pow1
+            return numpy.round(y) * pow1
+
+    return context.compile_internal(builder, round_ndigits, sig, args)
+
+@builtin
+@implement(numpy.round, types.Kind(types.Complex), types.Kind(types.Integer))
+def scalar_round_binary_complex(context, builder, sig, args):
+    def round_ndigits(z, ndigits):
+        return complex(numpy.round(z.real, ndigits),
+                       numpy.round(z.imag, ndigits))
+
+    return context.compile_internal(builder, round_ndigits, sig, args)
+
+
+@builtin
+@implement(numpy.round, types.Kind(types.Array), types.Kind(types.Integer),
+           types.Kind(types.Array))
+def array_round(context, builder, sig, args):
+    def array_round_impl(arr, decimals, out):
+        for x, y in zip(arr.shape, out.shape):
+            if x != y:
+                raise ValueError("invalid output shape")
+        for index, val in numpy.ndenumerate(arr):
+            out[index] = numpy.round(val, decimals)
+        return out
+
+    return context.compile_internal(builder, array_round_impl, sig, args)
+
+
 #-------------------------------------------------------------------------------
 
 
diff --git a/numba/typing/npydecl.py b/numba/typing/npydecl.py
--- a/numba/typing/npydecl.py
+++ b/numba/typing/npydecl.py
@@ -313,4 +313,39 @@ def generic(self, args, kws):
 builtin_global(numpy.ndindex, types.Function(NdIndex))
 
 
+@builtin
+class Round(AbstractTemplate):
+    key = numpy.round
+
+    def generic(self, args, kws):
+        assert not kws
+        assert 1 <= len(args) <= 3
+
+        arg = args[0]
+        if len(args) == 1:
+            decimals = types.int32
+            out = None
+        else:
+            decimals = args[1]
+            if len(args) == 2:
+                out = None
+            else:
+                out = args[2]
+
+        supported_scalars = (types.Integer, types.Float, types.Complex)
+        if isinstance(arg, supported_scalars):
+            assert out is None
+            return signature(arg, *args)
+        if (isinstance(arg, types.Array) and isinstance(arg.dtype, supported_scalars) and
+            isinstance(out, types.Array) and isinstance(out.dtype, supported_scalars) and
+            out.ndim == arg.ndim):
+            # arg can only be complex if out is complex too
+            if (not isinstance(arg.dtype, types.Complex)
+                or isinstance(out.dtype, types.Complex)):
+                return signature(out, *args)
+
+builtin_global(numpy.round, types.Function(Round))
+builtin_global(numpy.around, types.Function(Round))
+
+
 builtin_global(numpy, types.Module(numpy))


Test patch:
diff --git a/numba/tests/test_array_methods.py b/numba/tests/test_array_methods.py
--- a/numba/tests/test_array_methods.py
+++ b/numba/tests/test_array_methods.py
@@ -7,6 +7,7 @@
 from numba import unittest_support as unittest
 from numba import typeof, types
 from numba.compiler import compile_isolated
+from numba.numpy_support import as_dtype
 from .support import TestCase
 
 
@@ -42,6 +43,37 @@ def np_ndindex_array(arr):
     return s
 
 
+def np_around_array(arr, decimals, out):
+    np.around(arr, decimals, out)
+
+def np_around_binary(val, decimals):
+    return np.around(val, decimals)
+
+def np_around_unary(val):
+    return np.around(val)
+
+def np_round_array(arr, decimals, out):
+    np.round(arr, decimals, out)
+
+def np_round_binary(val, decimals):
+    return np.round(val, decimals)
+
+def np_round_unary(val):
+    return np.round(val)
+
+def _fixed_np_round(arr, decimals=0, out=None):
+    """
+    A slightly bugfixed version of np.round().
+    """
+    if out is not None and arr.dtype.kind == 'c':
+        # workaround for https://github.com/numpy/numpy/issues/5779
+        _fixed_np_round(arr.real, decimals, out.real)
+        _fixed_np_round(arr.imag, decimals, out.imag)
+        return out
+    else:
+        return np.round(arr, decimals, out)
+
+
 def array_sum(arr):
     return arr.sum()
 
@@ -391,6 +423,87 @@ def test_std_magnitude(self):
         self.check_aggregation_magnitude(array_std)
         self.check_aggregation_magnitude(array_std_global)
 
+    def check_round_array(self, pyfunc):
+        def check_round(cfunc, values, inty, outty, decimals):
+            # Create input and output arrays of the right type
+            arr = values.astype(as_dtype(inty))
+            out = np.zeros_like(arr).astype(as_dtype(outty))
+            pyout = out.copy()
+            _fixed_np_round(arr, decimals, pyout)
+            cfunc(arr, decimals, out)
+            np.testing.assert_allclose(out, pyout)
+            # Output shape mismatch
+            with self.assertRaises(ValueError) as raises:
+                cfunc(arr, decimals, out[1:])
+            self.assertEqual(str(raises.exception),
+                             "invalid output shape")
+
+        def check_types(argtypes, outtypes, values):
+            for inty, outty in product(argtypes, outtypes):
+                cres = compile_isolated(pyfunc,
+                                        (types.Array(inty, 1, 'A'),
+                                         types.int32,
+                                         types.Array(outty, 1, 'A')))
+                cfunc = cres.entry_point
+                check_round(cres.entry_point, values, inty, outty, 0)
+                check_round(cres.entry_point, values, inty, outty, 1)
+                if not isinstance(outty, types.Integer):
+                    check_round(cres.entry_point, values * 10, inty, outty, -1)
+                else:
+                    # Avoid Numpy bug when output is an int:
+                    # https://github.com/numpy/numpy/issues/5777
+                    pass
+
+        values = np.array([-3.0, -2.5, -2.25, -1.5, 1.5, 2.25, 2.5, 2.75])
+
+        argtypes = (types.float64, types.float32, types.int32)
+        check_types(argtypes, argtypes, values)
+
+        argtypes = (types.complex64, types.complex128)
+        check_types(argtypes, argtypes, values * (1 - 1j))
+
+    def test_round_array(self):
+        self.check_round_array(np_round_array)
+
+    def test_around_array(self):
+        self.check_round_array(np_around_array)
+
+    def check_round_scalar(self, unary_pyfunc, binary_pyfunc):
+        base_values = [-3.0, -2.5, -2.25, -1.5, 1.5, 2.25, 2.5, 2.75]
+        complex_values = [x * (1 - 1j) for x in base_values]
+        int_values = [int(x) for x in base_values]
+        argtypes = (types.float64, types.float32, types.int32,
+                    types.complex64, types.complex128)
+        argvalues = [base_values, base_values, int_values,
+                     complex_values, complex_values]
+
+        pyfunc = binary_pyfunc
+        for ty, values in zip(argtypes, argvalues):
+            cres = compile_isolated(pyfunc, (ty, types.int32))
+            cfunc = cres.entry_point
+            for decimals in (1, 0, -1):
+                for v in values:
+                    if decimals > 0:
+                        v *= 10
+                    expected = np.round(v, decimals)
+                    got = cfunc(v, decimals)
+                    self.assertPreciseEqual(got, expected)
+
+        pyfunc = unary_pyfunc
+        for ty, values in zip(argtypes, argvalues):
+            cres = compile_isolated(pyfunc, (ty,))
+            cfunc = cres.entry_point
+            for v in values:
+                expected = np.round(v)
+                got = cfunc(v)
+                self.assertPreciseEqual(got, expected)
+
+    def test_round_scalar(self):
+        self.check_round_scalar(np_round_unary, np_round_binary)
+
+    def test_around_scalar(self):
+        self.check_round_scalar(np_around_unary, np_around_binary)
+
 
 # These form a testing product where each of the combinations are tested
 reduction_funcs = [array_sum, array_sum_global,