[pypy-commit] pypy numpypy-out: branch to handle "out" arg, first stab at tests and implementation for reduce functions

mattip noreply at buildbot.pypy.org
Fri Feb 3 14:51:39 CET 2012


Author: mattip
Branch: numpypy-out
Changeset: r52072:620d218f4ddf
Date: 2012-02-03 15:45 +0200
http://bitbucket.org/pypy/pypy/changeset/620d218f4ddf/

Log:	branch to handle "out" arg, first stab at tests and implementation
	for reduce functions

diff --git a/pypy/module/micronumpy/interp_numarray.py b/pypy/module/micronumpy/interp_numarray.py
--- a/pypy/module/micronumpy/interp_numarray.py
+++ b/pypy/module/micronumpy/interp_numarray.py
@@ -131,13 +131,20 @@
     descr_rmod = _binop_right_impl("mod")
 
     def _reduce_ufunc_impl(ufunc_name, promote_to_largest=False):
-        def impl(self, space, w_axis=None):
+        def impl(self, space, w_axis=None, w_out=None):
             if space.is_w(w_axis, space.w_None):
                 axis = -1
             else:
                 axis = space.int_w(w_axis)
+            if space.is_w(w_out, space.w_None):
+                out = None
+            elif not isinstance(w_out, W_NDimArray):
+                raise OperationError(space.w_TypeError, space.wrap(
+                                                    'output must be an array'))
+            else:
+                out = w_out
             return getattr(interp_ufuncs.get(space), ufunc_name).reduce(space,
-                                        self, True, promote_to_largest, axis)
+                                        self, True, promote_to_largest, axis, False, out)
         return func_with_new_name(impl, "reduce_%s_impl" % ufunc_name)
 
     descr_sum = _reduce_ufunc_impl("add")
@@ -484,14 +491,14 @@
             )
         return w_result
 
-    def descr_mean(self, space, w_axis=None):
+    def descr_mean(self, space, w_axis=None, w_out=None):
         if space.is_w(w_axis, space.w_None):
             w_axis = space.wrap(-1)
             w_denom = space.wrap(self.size)
         else:
             dim = space.int_w(w_axis)
             w_denom = space.wrap(self.shape[dim])
-        return space.div(self.descr_sum_promote(space, w_axis), w_denom)
+        return space.div(self.descr_sum_promote(space, w_axis, w_out), w_denom)
 
     def descr_var(self, space, w_axis=None):
         return get_appbridge_cache(space).call_method(space, '_var', self,
diff --git a/pypy/module/micronumpy/interp_ufuncs.py b/pypy/module/micronumpy/interp_ufuncs.py
--- a/pypy/module/micronumpy/interp_ufuncs.py
+++ b/pypy/module/micronumpy/interp_ufuncs.py
@@ -105,28 +105,32 @@
         array([[ 1,  5],
                [ 9, 13]])
         """
-        if not space.is_w(w_out, space.w_None):
-            raise OperationError(space.w_NotImplementedError, space.wrap(
-                "out not supported"))
         if w_axis is None:
             axis = 0
         elif space.is_w(w_axis, space.w_None):
             axis = -1
         else:
             axis = space.int_w(w_axis)
-        return self.reduce(space, w_obj, False, False, axis, keepdims)
+        if space.is_w(w_out, space.w_None):
+            out = None
+        elif not isinstance(w_out, W_NDimArray):
+            raise OperationError(space.w_TypeError, space.wrap(
+                                                'output must be an array'))
+        else:
+            out = w_out
+        return self.reduce(space, w_obj, False, False, axis, keepdims, out)
 
-    def reduce(self, space, w_obj, multidim, promote_to_largest, dim,
-               keepdims=False):
+    def reduce(self, space, w_obj, multidim, promote_to_largest, axis,
+               keepdims=False, out=None):
         from pypy.module.micronumpy.interp_numarray import convert_to_array, \
-                                                           Scalar, ReduceArray
+                                             Scalar, ReduceArray, W_NDimArray
         if self.argcount != 2:
             raise OperationError(space.w_ValueError, space.wrap("reduce only "
                 "supported for binary functions"))
         assert isinstance(self, W_Ufunc2)
         obj = convert_to_array(space, w_obj)
-        if dim >= len(obj.shape):
-            raise OperationError(space.w_ValueError, space.wrap("axis(=%d) out of bounds" % dim))
+        if axis >= len(obj.shape):
+            raise OperationError(space.w_ValueError, space.wrap("axis(=%d) out of bounds" % axis))
         if isinstance(obj, Scalar):
             raise OperationError(space.w_TypeError, space.wrap("cannot reduce "
                 "on a scalar"))
@@ -144,21 +148,32 @@
         if self.identity is None and size == 0:
             raise operationerrfmt(space.w_ValueError, "zero-size array to "
                     "%s.reduce without identity", self.name)
-        if shapelen > 1 and dim >= 0:
-            return self.do_axis_reduce(obj, dtype, dim, keepdims)
+        if shapelen > 1 and axis >= 0:
+            if keepdims:
+                shape = obj.shape[:axis] + [1] + obj.shape[axis + 1:]
+            else:
+                shape = obj.shape[:axis] + obj.shape[axis + 1:]
+            if out:
+                #Test for shape agreement
+                return self.do_axis_reduce(obj, dtype, axis, out)
+            else:
+                result = W_NDimArray(support.product(shape), shape, dtype)
+                return self.do_axis_reduce(obj, dtype, axis, result)
         arr = ReduceArray(self.func, self.name, self.identity, obj, dtype)
-        return loop.compute(arr)
+        val = loop.compute(arr)
+        if out:
+            if len(out.shape)>0:
+                raise operationerrfmt(space.w_ValueError, "output parameter "
+                              "for reduction operation %s has too many"
+                              " dimensions",self.name)
+            out.setitem(0, out.dtype.coerce(space, val))
+            return out
+        return val 
 
-    def do_axis_reduce(self, obj, dtype, dim, keepdims):
-        from pypy.module.micronumpy.interp_numarray import AxisReduce,\
-             W_NDimArray
-        if keepdims:
-            shape = obj.shape[:dim] + [1] + obj.shape[dim + 1:]
-        else:
-            shape = obj.shape[:dim] + obj.shape[dim + 1:]
-        result = W_NDimArray(support.product(shape), shape, dtype)
+    def do_axis_reduce(self, obj, dtype, axis, result):
+        from pypy.module.micronumpy.interp_numarray import AxisReduce
         arr = AxisReduce(self.func, self.name, self.identity, obj.shape, dtype,
-                         result, obj, dim)
+                         result, obj, axis)
         loop.compute(arr)
         return arr.left
 
diff --git a/pypy/module/micronumpy/test/test_numarray.py b/pypy/module/micronumpy/test/test_numarray.py
--- a/pypy/module/micronumpy/test/test_numarray.py
+++ b/pypy/module/micronumpy/test/test_numarray.py
@@ -785,7 +785,7 @@
         assert (arange(10).reshape(5, 2).mean(axis=1) == [0.5, 2.5, 4.5, 6.5, 8.5]).all()
 
     def test_sum(self):
-        from _numpypy import array
+        from _numpypy import array,zeros
         a = array(range(5))
         assert a.sum() == 10
         assert a[:4].sum() == 6
@@ -794,6 +794,10 @@
         assert a.sum() == 5
 
         raises(TypeError, 'a.sum(2, 3)')
+        skip('fails since Scalar is not a subclass of W_NDimArray')
+        d = zeros(())
+        b = a.sum(out=d)
+        assert b == d
 
     def test_reduce_nd(self):
         from numpypy import arange, array, multiply
@@ -821,6 +825,14 @@
         assert (array([[1,2],[3,4]]).prod(0) == [3, 8]).all()
         assert (array([[1,2],[3,4]]).prod(1) == [2, 12]).all()
 
+    def test_reduce_out(self):
+        from numpypy import arange, array, multiply
+        a = arange(15).reshape(5, 3)
+        b = arange(3)
+        c = a.sum(0, out=b)
+        assert (c == [30, 35, 40]).all()
+        assert (c == b).all()
+
     def test_identity(self):
         from _numpypy import identity, array
         from _numpypy import int32, float64, dtype


More information about the pypy-commit mailing list