[pypy-commit] pypy ufuncapi: test broadcasting, add necessary but failing nditer itershape tests

[mattip]

Author: mattip <matti.picus at gmail.com>
Branch: ufuncapi
Changeset: r75293:a249b7e1b6ae
Date: 2015-01-10 23:55 +0200
http://bitbucket.org/pypy/pypy/changeset/a249b7e1b6ae/

Log:	test broadcasting, add necessary but failing nditer itershape tests

diff --git a/pypy/module/micronumpy/test/test_nditer.py b/pypy/module/micronumpy/test/test_nditer.py
--- a/pypy/module/micronumpy/test/test_nditer.py
+++ b/pypy/module/micronumpy/test/test_nditer.py
@@ -167,7 +167,8 @@
         exc = raises(TypeError, nditer, a, op_dtypes=['complex'])
         assert str(exc.value).startswith("Iterator operand required copying or buffering")
         exc = raises(ValueError, nditer, a, op_flags=['copy'], op_dtypes=['complex128'])
-        assert str(exc.value) == "None of the iterator flags READWRITE, READONLY, or WRITEONLY were specified for an operand"
+        assert str(exc.value) == "None of the iterator flags READWRITE," \
+                    " READONLY, or WRITEONLY were specified for an operand"
         r = []
         for x in nditer(a, op_flags=['readonly','copy'],
                         op_dtypes=['complex128']):
@@ -320,3 +321,34 @@
         assert res == [(0, (0, 0)), (1, (0, 1)),
                        (2, (0, 2)), (3, (1, 0)),
                        (4, (1, 1)), (5, (1, 2))]
+
+    def test_itershape(self):
+        # Check that allocated outputs work with a specified shape
+        from numpy import nditer, arange
+        a = arange(6, dtype='i2').reshape(2,3)
+        i = nditer([a, None], [], [['readonly'], ['writeonly','allocate']],
+                            op_axes=[[0,1,None], None],
+                            itershape=(-1,-1,4))
+        assert_equal(i.operands[1].shape, (2,3,4))
+        assert_equal(i.operands[1].strides, (24,8,2))
+
+        i = nditer([a.T, None], [], [['readonly'], ['writeonly','allocate']],
+                            op_axes=[[0,1,None], None],
+                            itershape=(-1,-1,4))
+        assert_equal(i.operands[1].shape, (3,2,4))
+        assert_equal(i.operands[1].strides, (8,24,2))
+
+        i = nditer([a.T, None], [], [['readonly'], ['writeonly','allocate']],
+                            order='F',
+                            op_axes=[[0,1,None], None],
+                            itershape=(-1,-1,4))
+        assert_equal(i.operands[1].shape, (3,2,4))
+        assert_equal(i.operands[1].strides, (2,6,12))
+
+        # If we specify 1 in the itershape, it shouldn't allow broadcasting
+        # of that dimension to a bigger value
+        assert_raises(ValueError, nditer, [a, None], [],
+                            [['readonly'], ['writeonly','allocate']],
+                            op_axes=[[0,1,None], None],
+                            itershape=(-1,1,4))
+
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
@@ -245,6 +245,31 @@
         ai2 = ufunc(ai)
         assert (ai2 == ai * 2).all()
 
+    def test_frompyfunc_sig_broadcast(self):
+        def sum_along_0(in_array, out_array):
+            out_array[...] = in_array.sum(axis=0)
+
+        def add_two(in0, in1, out):
+            out[...] = in0 + in1
+
+        from numpy import frompyfunc, dtype, arange
+        ufunc_add = frompyfunc(add_two, 2, 1,
+                            signature='(m,n),(m,n)->(m,n)',
+                            dtypes=[dtype(int), dtype(int), dtype(int)],
+                            stack_inputs=True,
+                          )
+        ufunc_sum = frompyfunc([sum_along_0], 1, 1,
+                            signature='(m,n)->(n)',
+                            dtypes=[dtype(int), dtype(int)],
+                            stack_inputs=True,
+                          )
+        ai = arange(18, dtype=int).reshape(3,2,3)
+        aout = ufunc_add(ai, ai[0,:,:])
+        assert aout.shape == (3, 2, 3)
+        aout = ufunc_sum(ai)
+        assert aout.shape == (3, 3)
+
+
     def test_ufunc_kwargs(self):
         from numpy import ufunc, frompyfunc, arange, dtype
         def adder(a, b):
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
@@ -869,7 +869,7 @@
                          self.name, name, _i, j, x, y)
                 iter_shape[offset + j] = max(x, y)
             #print 'Find or verify signature ixs',self.core_dim_ixs,
-            #print 'starting',dim_offset,'n',num_dims,'matching',dims_to_match
+            #print 'starting',dim_offset,'n',n,'num_dims',num_dims,'matching',dims_to_match
             for j in range(num_dims):
                 core_dim_index = self.core_dim_ixs[dim_offset + j]
                 if core_dim_index > len(dims_to_match):
@@ -886,7 +886,13 @@
                          self.name, name, _i, j, 
                          self.signature, matched_dims[core_dim_index],
                          dims_to_match[core_dim_index])
-            arg_shapes.append(iter_shape + dims_to_match)
+            #print 'adding',iter_shape,'+',dims_to_match,'to arg_shapes'
+            if n < len(iter_shape):
+                #Broadcast over the len(iter_shape) - n dims of iter_shape
+                broadcast_dims = len(iter_shape) - n
+                arg_shapes.append(iter_shape[:-broadcast_dims] + [1] * broadcast_dims + dims_to_match)
+            else:
+                arg_shapes.append(iter_shape + dims_to_match)
         # TODO once we support obejct dtypes,
         # FAIL with NotImplementedError if the other object has
         # the __r<op>__ method and has a higher priority than