Commits

mattip  committed fc93483

broadcast only when needed, fix compress for nd array with no axis arg (rguillebert, mattip)

  • Participants
  • Parent commits 628865d
  • Branches indexing-by-array

Comments (0)

Files changed (3)

File pypy/module/micronumpy/interp_numarray.py

             raise OperationError(space.w_ValueError,
                                  space.wrap("index out of range for array"))
         size = loop.count_all_true(arr)
-        print 'size',size
         if len(arr.get_shape()) == 1:
             res_shape = [size] + self.get_shape()[1:]
         else:
             res_shape = [size]
-        print 'res_shape',res_shape
         res = W_NDimArray.from_shape(res_shape, self.get_dtype())
         return loop.getitem_filter(res, self, arr)
 
         if not space.is_none(w_axis):
             raise OperationError(space.w_NotImplementedError,
                                  space.wrap("axis unsupported for compress"))
+            arr = self
+        else:
+            arr = self.descr_reshape(space, [space.wrap(-1)])
         index = convert_to_array(space, w_obj)
-        return self.getitem_filter(space, index)
+        return arr.getitem_filter(space, index)
 
     def descr_flatten(self, space, w_order=None):
         if self.is_scalar():

File pypy/module/micronumpy/loop.py

 
 def getitem_filter(res, arr, index):
     res_iter = res.create_iter()
-    print 'index.size', index.get_size()
-    print 'res.size', res.get_size()
-    index_iter = index.create_iter(arr.get_shape(), backward_broadcast=True)
+    shapelen = len(arr.get_shape())
+    if shapelen > 1 and len(index.get_shape()) < 2:
+        index_iter = index.create_iter(arr.get_shape(), backward_broadcast=True)
+    else:
+        index_iter = index.create_iter()
     arr_iter = arr.create_iter()
-    shapelen = len(arr.get_shape())
     arr_dtype = arr.get_dtype()
     index_dtype = index.get_dtype()
     # XXX length of shape of index as well?
                                               index_dtype=index_dtype,
                                               arr_dtype=arr_dtype,
                                               )
-        print 'res,arr,index', res_iter.offset, arr_iter.offset, index_iter.offset, index_iter.getitem_bool()
         if index_iter.getitem_bool():
             res_iter.setitem(arr_iter.getitem())
             res_iter.next()

File pypy/module/micronumpy/test/test_numarray.py

     def test_compress(self):
         from numpypy import arange, array
         a = arange(10)
-        print 0
         assert (a.compress([True, False, True]) == [0, 2]).all()
-        print 1
         assert (a.compress([1, 0, 13]) == [0, 2]).all()
-        print 2
         assert (a.compress([1, 0, 13]) == [0, 2]).all()
-        print '2a'
         assert (a.compress([1, 0, 13.5]) == [0, 2]).all()
-        print 3
         assert (a.compress(array([1, 0, 13.5], dtype='>f4')) == [0, 2]).all()
-        print 4
         assert (a.compress(array([1, 0, 13.5], dtype='<f4')) == [0, 2]).all()
-        print 5
         assert (a.compress([1, -0-0j, 1.3+13.5j]) == [0, 2]).all()
-        print 6
         a = arange(10).reshape(2, 5)
         assert (a.compress([True, False, True]) == [0, 2]).all()
-        print 7
         raises((IndexError, ValueError), "a.compress([1] * 100)")
-        print 8
 
     def test_item(self):
         from numpypy import array