[pypy-commit] pypy ufunc-reduce: Replace 'cumulative' flag with 'variant' enum in reduce()

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: ufunc-reduce
Changeset: r78643:5cb612e5ab3a
Date: 2015-07-23 19:59 +0100
http://bitbucket.org/pypy/pypy/changeset/5cb612e5ab3a/

Log:	Replace 'cumulative' flag with 'variant' enum in reduce()

diff --git a/pypy/module/micronumpy/ndarray.py b/pypy/module/micronumpy/ndarray.py
--- a/pypy/module/micronumpy/ndarray.py
+++ b/pypy/module/micronumpy/ndarray.py
@@ -1149,7 +1149,7 @@
 
     # ----------------------- reduce -------------------------------
 
-    def _reduce_ufunc_impl(ufunc_name, cumulative=False, bool_result=False):
+    def _reduce_ufunc_impl(ufunc_name, name, variant=ufuncs.REDUCE, bool_result=False):
         @unwrap_spec(keepdims=bool)
         def impl(self, space, w_axis=None, w_dtype=None, w_out=None, keepdims=False):
             if space.is_none(w_out):
@@ -1161,18 +1161,19 @@
             if bool_result:
                 w_dtype = descriptor.get_dtype_cache(space).w_booldtype
             return getattr(ufuncs.get(space), ufunc_name).reduce(
-                space, self, w_axis, keepdims, out, w_dtype, cumulative=cumulative)
-        return func_with_new_name(impl, "reduce_%s_impl_%d" % (ufunc_name, cumulative))
+                space, self, w_axis, keepdims, out, w_dtype, variant=variant)
+        impl.__name__ = name
+        return impl
 
-    descr_sum = _reduce_ufunc_impl("add")
-    descr_prod = _reduce_ufunc_impl("multiply")
-    descr_max = _reduce_ufunc_impl("maximum")
-    descr_min = _reduce_ufunc_impl("minimum")
-    descr_all = _reduce_ufunc_impl('logical_and', bool_result=True)
-    descr_any = _reduce_ufunc_impl('logical_or', bool_result=True)
+    descr_sum = _reduce_ufunc_impl("add", "descr_sum")
+    descr_prod = _reduce_ufunc_impl("multiply", "descr_prod")
+    descr_max = _reduce_ufunc_impl("maximum", "descr_max")
+    descr_min = _reduce_ufunc_impl("minimum", "descr_min")
+    descr_all = _reduce_ufunc_impl('logical_and', "descr_all", bool_result=True)
+    descr_any = _reduce_ufunc_impl('logical_or', "descr_any", bool_result=True)
 
-    descr_cumsum = _reduce_ufunc_impl('add', cumulative=True)
-    descr_cumprod = _reduce_ufunc_impl('multiply', cumulative=True)
+    descr_cumsum = _reduce_ufunc_impl('add', "descr_cumsum", variant=ufuncs.ACCUMULATE)
+    descr_cumprod = _reduce_ufunc_impl('multiply', "descr_cumprod", variant=ufuncs.ACCUMULATE)
 
     def _reduce_argmax_argmin_impl(raw_name):
         op_name = "arg%s" % raw_name
diff --git a/pypy/module/micronumpy/ufuncs.py b/pypy/module/micronumpy/ufuncs.py
--- a/pypy/module/micronumpy/ufuncs.py
+++ b/pypy/module/micronumpy/ufuncs.py
@@ -25,6 +25,8 @@
     find_result_type, promote_types)
 from .boxes import W_GenericBox, W_ObjectBox
 
+REDUCE, ACCUMULATE, REDUCEAT = range(3)
+
 def done_if_true(dtype, val):
     return dtype.itemtype.bool(val)
 
@@ -171,7 +173,7 @@
         else:
             out = w_out
         return self.reduce(space, w_obj, w_axis, True, #keepdims must be true
-                           out, w_dtype, cumulative=True)
+                           out, w_dtype, variant=ACCUMULATE)
 
     @unwrap_spec(keepdims=bool)
     def descr_reduce(self, space, w_obj, w_axis=None, w_dtype=None,
@@ -241,7 +243,7 @@
         return self.reduce(space, w_obj, w_axis, keepdims, out, w_dtype)
 
     def reduce(self, space, w_obj, w_axis, keepdims=False, out=None, dtype=None,
-               cumulative=False):
+               variant=REDUCE):
         if self.nin != 2:
             raise oefmt(space.w_ValueError,
                         "reduce only supported for binary functions")
@@ -292,7 +294,7 @@
                             "zero-size array to reduction operation %s "
                             "which has no identity", self.name)
 
-        if cumulative:
+        if variant == ACCUMULATE:
             dtype = self.find_binop_type(space, dtype)
         else:
             _, dtype, _ = self.find_specialization(space, dtype, dtype, out,
@@ -300,7 +302,7 @@
         call__array_wrap__ = True
         if shapelen > 1 and axis < shapelen:
             temp = None
-            if cumulative:
+            if variant == ACCUMULATE:
                 shape = obj_shape[:]
                 temp_shape = obj_shape[:axis] + obj_shape[axis + 1:]
                 if out:
@@ -339,13 +341,13 @@
                 if self.identity is not None:
                     out.fill(space, self.identity.convert_to(space, dtype))
                 return out
-            loop.do_axis_reduce(space, shape, self.func, obj, dtype,
-                                       axis, out, self.identity, cumulative,
-                                       temp)
+            loop.do_axis_reduce(space, shape, self.func, obj, dtype, axis,
+                                out, self.identity, (variant == ACCUMULATE),
+                                temp)
             if call__array_wrap__:
                 out = space.call_method(obj, '__array_wrap__', out)
             return out
-        if cumulative:
+        if variant == ACCUMULATE:
             if out:
                 call__array_wrap__ = False
                 if out.get_shape() != [obj.get_size()]:
@@ -786,7 +788,7 @@
         self.external_loop = external_loop
 
     def reduce(self, space, w_obj, w_axis, keepdims=False, out=None, dtype=None,
-               cumulative=False):
+               variant=REDUCE):
         raise oefmt(space.w_NotImplementedError, 'not implemented yet')
 
     def call(self, space, args_w, sig, casting, extobj):