[pypy-commit] pypy ufunc-reduce: Support multiple axes in ufunc.reduce()

[rlamy]

Author: Ronan Lamy <ronan.lamy at gmail.com>
Branch: ufunc-reduce
Changeset: r78699:cfb5865ce1e9
Date: 2015-07-28 19:19 +0100
http://bitbucket.org/pypy/pypy/changeset/cfb5865ce1e9/

Log:	Support multiple axes in ufunc.reduce()

diff --git a/pypy/module/micronumpy/test/test_ufuncs.py b/pypy/module/micronumpy/test/test_ufuncs.py
--- a/pypy/module/micronumpy/test/test_ufuncs.py
+++ b/pypy/module/micronumpy/test/test_ufuncs.py
@@ -1044,6 +1044,13 @@
             assert np.equal.reduce([1, 2], dtype=dtype) == True
             assert np.equal.reduce([1, 2, 0], dtype=dtype) == False
 
+    def test_reduce_axes(self):
+        import numpy as np
+        a = np.arange(24).reshape(2, 3, 4)
+        b = np.add.reduce(a, axis=(0, 1))
+        assert b.shape == (4,)
+        assert (b == [60, 66, 72, 78]).all()
+
     def test_reduce_fmax(self):
         import numpy as np
         assert np.fmax.reduce(np.arange(11).astype('b')) == 10
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
@@ -247,6 +247,18 @@
         if space.is_none(w_axis):
             axes = range(shapelen)
             axis = maxint
+        elif space.isinstance_w(w_axis, space.w_tuple):
+            axes_w = space.listview(w_axis)
+            axes = [0] * len(axes_w)
+            for i in range(len(axes_w)):
+                x = space.int_w(axes_w[i])
+                if x < 0:
+                    x += shapelen
+                if x < 0 or x >= shapelen:
+                    raise oefmt(space.w_ValueError, "'axis' entry is out of bounds")
+                axes[i] = x
+
+
         else:
             if space.isinstance_w(w_axis, space.w_tuple) and space.len_w(w_axis) == 1:
                 w_axis = space.getitem(w_axis, space.wrap(0))
@@ -256,7 +268,6 @@
             if axis < 0:
                 axis += shapelen
             axes = [axis]
-        assert axis >= 0
         dtype = decode_w_dtype(space, dtype)
 
         if dtype is None and out is not None:
@@ -277,12 +288,11 @@
             dtype = num2dtype(space, num)
 
         if self.identity is None:
-            for i in range(shapelen):
-                if space.is_none(w_axis) or i == axis:
-                    if obj_shape[i] == 0:
-                        raise oefmt(space.w_ValueError,
-                            "zero-size array to reduction operation %s "
-                            "which has no identity", self.name)
+            for i in axes:
+                if obj_shape[i] == 0:
+                    raise oefmt(space.w_ValueError,
+                        "zero-size array to reduction operation %s "
+                        "which has no identity", self.name)
 
         if variant == ACCUMULATE:
             if len(axes) != 1:
@@ -379,9 +389,16 @@
         else:
             temp = None
             if keepdims:
-                shape = obj_shape[:axis] + [1] + obj_shape[axis + 1:]
+                shape = obj_shape[:]
+                for axis in axes:
+                    shape[axis] = 1
             else:
-                shape = obj_shape[:axis] + obj_shape[axis + 1:]
+                shape = [0] * (shapelen - len(axes))
+                j = 0
+                for i in range(shapelen):
+                    if not axis_flags[i]:
+                        shape[j] = obj_shape[i]
+                        j += 1
             if out:
                 # Test for shape agreement
                 # XXX maybe we need to do broadcasting here, although I must