[pypy-commit] pypy numpy-record-dtypes: break the world and then fix it. strides are now based on bytes not on

fijal noreply at buildbot.pypy.org
Fri Feb 17 14:42:35 CET 2012


Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: numpy-record-dtypes
Changeset: r52584:de511ede85c1
Date: 2012-02-17 15:42 +0200
http://bitbucket.org/pypy/pypy/changeset/de511ede85c1/

Log:	break the world and then fix it. strides are now based on bytes not
	on items. saves a mul in test_zjit as well

diff --git a/pypy/module/micronumpy/interp_dtype.py b/pypy/module/micronumpy/interp_dtype.py
--- a/pypy/module/micronumpy/interp_dtype.py
+++ b/pypy/module/micronumpy/interp_dtype.py
@@ -42,17 +42,16 @@
         return self.itemtype.coerce(space, self, w_item)
 
     def getitem(self, arr, i):
-        return self.itemtype.read(arr, self.itemtype.get_element_size(), i, 0)
+        return self.itemtype.read(arr, 1, i, 0)
 
     def getitem_bool(self, arr, i):
-        isize = self.itemtype.get_element_size()
-        return self.itemtype.read_bool(arr, isize, i, 0)
+        return self.itemtype.read_bool(arr, 1, i, 0)
 
     def setitem(self, arr, i, box):
-        self.itemtype.store(arr, self.itemtype.get_element_size(), i, 0, box)
+        self.itemtype.store(arr, 1, i, 0, box)
 
     def fill(self, storage, box, start, stop):
-        self.itemtype.fill(storage, self.itemtype.get_element_size(), box, start, stop, 0)
+        self.itemtype.fill(storage, self.get_size(), box, start, stop, 0)
 
     def descr_str(self, space):
         return space.wrap(self.name)
@@ -120,6 +119,9 @@
             return '<DType %r>' % self.fields
         return '<DType %r>' % self.itemtype
 
+    def get_size(self):
+        return self.itemtype.get_element_size()
+
 def dtype_from_list(space, w_lst):
     lst_w = space.listview(w_lst)
     fields = {}
diff --git a/pypy/module/micronumpy/interp_iter.py b/pypy/module/micronumpy/interp_iter.py
--- a/pypy/module/micronumpy/interp_iter.py
+++ b/pypy/module/micronumpy/interp_iter.py
@@ -42,7 +42,7 @@
 we can go faster.
 All the calculations happen in next()
 
-next_step_x() tries to do the iteration for a number of steps at once,
+next_skip_x() tries to do the iteration for a number of steps at once,
 but then we cannot gaurentee that we only overflow one single shape 
 dimension, perhaps we could overflow times in one big step.
 """
@@ -95,17 +95,19 @@
         raise NotImplementedError
 
 class ArrayIterator(BaseIterator):
-    def __init__(self, size):
+    def __init__(self, size, element_size):
         self.offset = 0
         self.size = size
+        self.element_size = element_size
 
     def next(self, shapelen):
         return self.next_skip_x(1)
 
-    def next_skip_x(self, ofs):
+    def next_skip_x(self, x):
         arr = instantiate(ArrayIterator)
         arr.size = self.size
-        arr.offset = self.offset + ofs
+        arr.offset = self.offset + x * self.element_size
+        arr.element_size = self.element_size
         return arr
 
     def next_no_increase(self, shapelen):
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
@@ -79,8 +79,8 @@
         dtype = space.interp_w(interp_dtype.W_Dtype,
             space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype)
         )
-        size, shape = _find_size_and_shape(space, w_size)
-        return space.wrap(W_NDimArray(size, shape[:], dtype=dtype))
+        shape = _find_shape(space, w_size)
+        return space.wrap(W_NDimArray(shape[:], dtype=dtype))
 
     def _unaryop_impl(ufunc_name):
         def impl(self, space):
@@ -204,8 +204,7 @@
             return scalar_w(space, dtype, space.wrap(0))
         # Do the dims match?
         out_shape, other_critical_dim = match_dot_shapes(space, self, other)
-        out_size = support.product(out_shape)
-        result = W_NDimArray(out_size, out_shape, dtype)
+        result = W_NDimArray(out_shape, dtype)
         # This is the place to add fpypy and blas
         return multidim_dot(space, self.get_concrete(),
                             other.get_concrete(), result, dtype,
@@ -224,7 +223,7 @@
         return space.wrap(self.find_dtype().itemtype.get_element_size())
 
     def descr_get_nbytes(self, space):
-        return space.wrap(self.size * self.find_dtype().itemtype.get_element_size())
+        return space.wrap(self.size)
 
     @jit.unroll_safe
     def descr_get_shape(self, space):
@@ -232,13 +231,16 @@
 
     def descr_set_shape(self, space, w_iterable):
         new_shape = get_shape_from_iterable(space,
-                            self.size, w_iterable)
+                            support.product(self.shape), w_iterable)
         if isinstance(self, Scalar):
             return
         self.get_concrete().setshape(space, new_shape)
 
     def descr_get_size(self, space):
-        return space.wrap(self.size)
+        return space.wrap(self.get_size())
+
+    def get_size(self):
+        return self.size // self.find_dtype().get_size()
 
     def descr_copy(self, space):
         return self.copy(space)
@@ -258,7 +260,7 @@
 
     def empty_copy(self, space, dtype):
         shape = self.shape
-        return W_NDimArray(support.product(shape), shape[:], dtype, 'C')
+        return W_NDimArray(shape[:], dtype, 'C')
 
     def descr_len(self, space):
         if len(self.shape):
@@ -299,6 +301,8 @@
         """ The result of getitem/setitem is a single item if w_idx
         is a list of scalars that match the size of shape
         """
+        if space.isinstance_w(w_idx, space.w_str):
+            return False
         shape_len = len(self.shape)
         if shape_len == 0:
             raise OperationError(space.w_IndexError, space.wrap(
@@ -330,28 +334,28 @@
         return [Chunk(*space.decode_index4(w_item, self.shape[i])) for i, w_item in
                 enumerate(space.fixedview(w_idx))]
 
-    def count_all_true(self, arr):
-        sig = arr.find_sig()
-        frame = sig.create_frame(arr)
-        shapelen = len(arr.shape)
+    def count_all_true(self):
+        sig = self.find_sig()
+        frame = sig.create_frame(self)
+        shapelen = len(self.shape)
         s = 0
         iter = None
         while not frame.done():
-            count_driver.jit_merge_point(arr=arr, frame=frame, iter=iter, s=s,
+            count_driver.jit_merge_point(arr=self, frame=frame, iter=iter, s=s,
                                          shapelen=shapelen)
             iter = frame.get_final_iter()
-            s += arr.dtype.getitem_bool(arr, iter.offset)
+            s += self.dtype.getitem_bool(self, iter.offset)
             frame.next(shapelen)
         return s
 
     def getitem_filter(self, space, arr):
         concr = arr.get_concrete()
-        if concr.size > self.size:
+        if concr.get_size() > self.get_size():
             raise OperationError(space.w_IndexError,
                                  space.wrap("index out of range for array"))
-        size = self.count_all_true(concr)
-        res = W_NDimArray(size, [size], self.find_dtype())
-        ri = ArrayIterator(size)
+        size = concr.count_all_true()
+        res = W_NDimArray([size], self.find_dtype())
+        ri = res.create_iter()
         shapelen = len(self.shape)
         argi = concr.create_iter()
         sig = self.find_sig()
@@ -372,7 +376,7 @@
         return res
 
     def setitem_filter(self, space, idx, val):
-        size = self.count_all_true(idx)
+        size = idx.count_all_true()
         arr = SliceArray([size], self.dtype, self, val)
         sig = arr.find_sig()
         shapelen = len(self.shape)
@@ -451,12 +455,13 @@
             w_shape = args_w[0]
         else:
             w_shape = space.newtuple(args_w)
-        new_shape = get_shape_from_iterable(space, self.size, w_shape)
+        new_shape = get_shape_from_iterable(space, support.product(self.shape),
+                                            w_shape)
         return self.reshape(space, new_shape)
 
     def reshape(self, space, new_shape):
         concrete = self.get_concrete()
-        # Since we got to here, prod(new_shape) == self.size
+        # Since we got to here, prod(new_shape) == self.get_size()
         new_strides = calc_new_strides(new_shape, concrete.shape,
                                      concrete.strides, concrete.order)
         if new_strides:
@@ -487,7 +492,7 @@
     def descr_mean(self, space, w_axis=None):
         if space.is_w(w_axis, space.w_None):
             w_axis = space.wrap(-1)
-            w_denom = space.wrap(self.size)
+            w_denom = space.wrap(support.product(self.shape))
         else:
             dim = space.int_w(w_axis)
             w_denom = space.wrap(self.shape[dim])
@@ -506,7 +511,7 @@
         concr.fill(space, w_value)
 
     def descr_nonzero(self, space):
-        if self.size > 1:
+        if self.get_size() > 1:
             raise OperationError(space.w_ValueError, space.wrap(
                 "The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()"))
         concr = self.get_concrete_or_scalar()
@@ -585,8 +590,7 @@
                                  space.wrap("axis unsupported for take"))
         index_i = index.create_iter()
         res_shape = index.shape
-        size = support.product(res_shape)
-        res = W_NDimArray(size, res_shape[:], concr.dtype, concr.order)
+        res = W_NDimArray(res_shape[:], concr.dtype, concr.order)
         res_i = res.create_iter()
         shapelen = len(index.shape)
         sig = concr.find_sig()
@@ -651,8 +655,7 @@
     """
     Intermediate class representing a literal.
     """
-    size = 1
-    _attrs_ = ["dtype", "value", "shape"]
+    _attrs_ = ["dtype", "value", "shape", "size"]
 
     def __init__(self, dtype, value):
         self.shape = []
@@ -660,6 +663,7 @@
         self.dtype = dtype
         assert isinstance(value, interp_boxes.W_GenericBox)
         self.value = value
+        self.size = dtype.get_size()
 
     def find_dtype(self):
         return self.dtype
@@ -677,8 +681,7 @@
         return self
 
     def reshape(self, space, new_shape):
-        size = support.product(new_shape)
-        res = W_NDimArray(size, new_shape, self.dtype, 'C')
+        res = W_NDimArray(new_shape, self.dtype, 'C')
         res.setitem(0, self.value)
         return res
 
@@ -691,6 +694,7 @@
         self.forced_result = None
         self.res_dtype = res_dtype
         self.name = name
+        self.size = support.product(self.shape) * res_dtype.get_size()
 
     def _del_sources(self):
         # Function for deleting references to source arrays,
@@ -698,7 +702,7 @@
         raise NotImplementedError
 
     def compute(self):
-        ra = ResultArray(self, self.size, self.shape, self.res_dtype)
+        ra = ResultArray(self, self.shape, self.res_dtype)
         loop.compute(ra)
         return ra.left
 
@@ -726,7 +730,6 @@
     def __init__(self, child, chunks, shape):
         self.child = child
         self.chunks = chunks
-        self.size = support.product(shape)
         VirtualArray.__init__(self, 'slice', shape, child.find_dtype())
 
     def create_sig(self):
@@ -773,7 +776,6 @@
         self.left = left
         self.right = right
         self.calc_dtype = calc_dtype
-        self.size = support.product(self.shape)
 
     def _del_sources(self):
         self.left = None
@@ -801,9 +803,9 @@
                                self.left.create_sig(), self.right.create_sig())
 
 class ResultArray(Call2):
-    def __init__(self, child, size, shape, dtype, res=None, order='C'):
+    def __init__(self, child, shape, dtype, res=None, order='C'):
         if res is None:
-            res = W_NDimArray(size, shape, dtype, order)
+            res = W_NDimArray(shape, dtype, order)
         Call2.__init__(self, None, 'assign', shape, dtype, dtype, res, child)
 
     def create_sig(self):
@@ -817,7 +819,7 @@
         self.s = StringBuilder(child.size * self.itemsize)
         Call1.__init__(self, None, 'tostring', child.shape, dtype, dtype,
                        child)
-        self.res = W_NDimArray(1, [1], dtype, 'C')
+        self.res = W_NDimArray([1], dtype, 'C')
         self.res_casted = rffi.cast(rffi.CArrayPtr(lltype.Char),
                                     self.res.storage)
 
@@ -898,13 +900,13 @@
     """
     _immutable_fields_ = ['storage']
 
-    def __init__(self, size, shape, dtype, order='C', parent=None):
-        self.size = size
+    def __init__(self, shape, dtype, order='C', parent=None):
         self.parent = parent
+        self.size = support.product(shape) * dtype.get_size()
         if parent is not None:
             self.storage = parent.storage
         else:
-            self.storage = dtype.itemtype.malloc(size)
+            self.storage = dtype.itemtype.malloc(self.size)
         self.order = order
         self.dtype = dtype
         if self.strides is None:
@@ -930,6 +932,7 @@
         self.dtype.setitem(self, item, value)
 
     def calc_strides(self, shape):
+        dtype = self.find_dtype()
         strides = []
         backstrides = []
         s = 1
@@ -937,8 +940,8 @@
         if self.order == 'C':
             shape_rev.reverse()
         for sh in shape_rev:
-            strides.append(s)
-            backstrides.append(s * (sh - 1))
+            strides.append(s * dtype.get_size())
+            backstrides.append(s * (sh - 1) * dtype.get_size())
             s *= sh
         if self.order == 'C':
             strides.reverse()
@@ -986,9 +989,9 @@
         shapelen = len(self.shape)
         if shapelen == 1:
             rffi.c_memcpy(
-                rffi.ptradd(self.storage, self.start * itemsize),
-                rffi.ptradd(w_value.storage, w_value.start * itemsize),
-                self.size * itemsize
+                rffi.ptradd(self.storage, self.start),
+                rffi.ptradd(w_value.storage, w_value.start),
+                self.size
             )
         else:
             dest = SkipLastAxisIterator(self)
@@ -1003,7 +1006,7 @@
                 dest.next()
 
     def copy(self, space):
-        array = W_NDimArray(self.size, self.shape[:], self.dtype, self.order)
+        array = W_NDimArray(self.shape[:], self.dtype, self.order)
         array.setslice(space, self)
         return array
 
@@ -1023,8 +1026,7 @@
             parent = parent.parent
         self.strides = strides
         self.backstrides = backstrides
-        ViewArray.__init__(self, support.product(shape), shape, parent.dtype,
-                           parent.order, parent)
+        ViewArray.__init__(self, shape, parent.dtype, parent.order, parent)
         self.start = start
 
     def create_iter(self, transforms=None):
@@ -1039,12 +1041,13 @@
             # but then calc_strides would have to accept a stepping factor
             strides = []
             backstrides = []
-            s = self.strides[0]
+            dtype = self.find_dtype()
+            s = self.strides[0] // dtype.get_size()
             if self.order == 'C':
                 new_shape.reverse()
             for sh in new_shape:
-                strides.append(s)
-                backstrides.append(s * (sh - 1))
+                strides.append(s * dtype.get_size())
+                backstrides.append(s * (sh - 1) * dtype.get_size())
                 s *= sh
             if self.order == 'C':
                 strides.reverse()
@@ -1079,7 +1082,9 @@
         self.calc_strides(new_shape)
 
     def create_iter(self, transforms=None):
-        return ArrayIterator(self.size).apply_transformations(self, transforms)
+        esize = self.find_dtype().get_size()
+        return ArrayIterator(self.size, esize).apply_transformations(self,
+                                                                     transforms)
 
     def create_sig(self):
         return signature.ArraySignature(self.dtype)
@@ -1087,18 +1092,13 @@
     def __del__(self):
         lltype.free(self.storage, flavor='raw', track_allocation=False)
 
-def _find_size_and_shape(space, w_size):
+def _find_shape(space, w_size):
     if space.isinstance_w(w_size, space.w_int):
-        size = space.int_w(w_size)
-        shape = [size]
-    else:
-        size = 1
-        shape = []
-        for w_item in space.fixedview(w_size):
-            item = space.int_w(w_item)
-            size *= item
-            shape.append(item)
-    return size, shape
+        return [space.int_w(w_size)]
+    shape = []
+    for w_item in space.fixedview(w_size):
+        shape.append(space.int_w(w_item))
+    return shape
 
 @unwrap_spec(subok=bool, copy=bool, ownmaskna=bool)
 def array(space, w_item_or_iterable, w_dtype=None, w_order=None,
@@ -1139,16 +1139,17 @@
            space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype))
     shape, elems_w = find_shape_and_elems(space, w_item_or_iterable, dtype)
     # they come back in C order
-    size = len(elems_w)
     if dtype is None:
         for w_elem in elems_w:
             dtype = interp_ufuncs.find_dtype_for_scalar(space, w_elem,
-                                                          w_dtype)
+                                                        dtype)
             if dtype is interp_dtype.get_dtype_cache(space).w_float64dtype:
                 break
-    arr = W_NDimArray(size, shape[:], dtype=dtype, order=order)
+        if dtype is None:
+            dtype = interp_dtype.get_dtype_cache(space).w_float64dtype
+    arr = W_NDimArray(shape[:], dtype=dtype, order=order)
     shapelen = len(shape)
-    arr_iter = ArrayIterator(arr.size)
+    arr_iter = arr.create_iter()
     # XXX we might want to have a jitdriver here
     for i in range(len(elems_w)):
         w_elem = elems_w[i]
@@ -1161,22 +1162,22 @@
     dtype = space.interp_w(interp_dtype.W_Dtype,
         space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype)
     )
-    size, shape = _find_size_and_shape(space, w_size)
+    shape = _find_shape(space, w_size)
     if not shape:
         return scalar_w(space, dtype, space.wrap(0))
-    return space.wrap(W_NDimArray(size, shape[:], dtype=dtype))
+    return space.wrap(W_NDimArray(shape[:], dtype=dtype))
 
 def ones(space, w_size, w_dtype=None):
     dtype = space.interp_w(interp_dtype.W_Dtype,
         space.call_function(space.gettypefor(interp_dtype.W_Dtype), w_dtype)
     )
 
-    size, shape = _find_size_and_shape(space, w_size)
+    shape = _find_shape(space, w_size)
     if not shape:
         return scalar_w(space, dtype, space.wrap(1))
-    arr = W_NDimArray(size, shape[:], dtype=dtype)
+    arr = W_NDimArray(shape[:], dtype=dtype)
     one = dtype.box(1)
-    arr.dtype.fill(arr.storage, one, 0, size)
+    arr.dtype.fill(arr.storage, one, 0, arr.size)
     return space.wrap(arr)
 
 @unwrap_spec(arr=BaseArray, skipna=bool, keepdims=bool)
@@ -1224,7 +1225,7 @@
                     "array dimensions must agree except for axis being concatenated"))
             elif i == axis:
                 shape[i] += axis_size
-    res = W_NDimArray(support.product(shape), shape, dtype, 'C')
+    res = W_NDimArray(shape, dtype, 'C')
     chunks = [Chunk(0, i, 1, i) for i in shape]
     axis_start = 0
     for arr in args_w:
@@ -1327,7 +1328,7 @@
         self.iter = sig.create_frame(arr).get_final_iter()
         self.base = arr
         self.index = 0
-        ViewArray.__init__(self, arr.size, [arr.size], arr.dtype, arr.order,
+        ViewArray.__init__(self, [arr.get_size()], arr.dtype, arr.order,
                            arr)
 
     def descr_next(self, space):
@@ -1342,7 +1343,7 @@
         return self
 
     def descr_len(self, space):
-        return space.wrap(self.size)
+        return space.wrap(self.get_size())
 
     def descr_index(self, space):
         return space.wrap(self.index)
@@ -1360,18 +1361,17 @@
             raise OperationError(space.w_IndexError,
                                  space.wrap('unsupported iterator index'))
         base = self.base
-        start, stop, step, lngth = space.decode_index4(w_idx, base.size)
+        start, stop, step, lngth = space.decode_index4(w_idx, base.get_size())
         # setslice would have been better, but flat[u:v] for arbitrary
         # shapes of array a cannot be represented as a[x1:x2, y1:y2]
         basei = ViewIterator(base.start, base.strides,
-                               base.backstrides,base.shape)
+                             base.backstrides, base.shape)
         shapelen = len(base.shape)
         basei = basei.next_skip_x(shapelen, start)
         if lngth <2:
             return base.getitem(basei.offset)
-        ri = ArrayIterator(lngth)
-        res = W_NDimArray(lngth, [lngth], base.dtype,
-                                    base.order)
+        res = W_NDimArray([lngth], base.dtype, base.order)
+        ri = res.create_iter()
         while not ri.done():
             flat_get_driver.jit_merge_point(shapelen=shapelen,
                                              base=base,
@@ -1381,7 +1381,7 @@
                                              ri=ri,
                                             )
             w_val = base.getitem(basei.offset)
-            res.setitem(ri.offset,w_val)
+            res.setitem(ri.offset, w_val)
             basei = basei.next_skip_x(shapelen, step)
             ri = ri.next(shapelen)
         return res
@@ -1392,11 +1392,12 @@
             raise OperationError(space.w_IndexError,
                                  space.wrap('unsupported iterator index'))
         base = self.base
-        start, stop, step, lngth = space.decode_index4(w_idx, base.size)
+        start, stop, step, lngth = space.decode_index4(w_idx, base.get_size())
         arr = convert_to_array(space, w_value)
+        ri = arr.create_iter()
         ai = 0
         basei = ViewIterator(base.start, base.strides,
-                               base.backstrides,base.shape)
+                             base.backstrides, base.shape)
         shapelen = len(base.shape)
         basei = basei.next_skip_x(shapelen, start)
         while lngth > 0:
@@ -1408,11 +1409,13 @@
                                              ai=ai,
                                              lngth=lngth,
                                             )
-            v = arr.getitem(ai).convert_to(base.dtype)
+            v = arr.getitem(ri.offset).convert_to(base.dtype)
             base.setitem(basei.offset, v)
             # need to repeat input values until all assignments are done
-            ai = (ai + 1) % arr.size
             basei = basei.next_skip_x(shapelen, step)
+            ri = ri.next(shapelen)
+            # WTF is numpy thinking?
+            ri.offset %= arr.size
             lngth -= 1
 
     def create_sig(self):
@@ -1420,9 +1423,9 @@
 
     def create_iter(self, transforms=None):
         return ViewIterator(self.base.start, self.base.strides,
-                    self.base.backstrides,
-                    self.base.shape).apply_transformations(self.base,
-                                                           transforms)
+                            self.base.backstrides,
+                            self.base.shape).apply_transformations(self.base,
+                                                                   transforms)
 
     def descr_base(self, space):
         return space.wrap(self.base)
diff --git a/pypy/module/micronumpy/interp_support.py b/pypy/module/micronumpy/interp_support.py
--- a/pypy/module/micronumpy/interp_support.py
+++ b/pypy/module/micronumpy/interp_support.py
@@ -51,9 +51,11 @@
         raise OperationError(space.w_ValueError, space.wrap(
             "string is smaller than requested size"))
 
-    a = W_NDimArray(num_items, [num_items], dtype=dtype)
-    for i, val in enumerate(items):
-        a.dtype.setitem(a, i, val)
+    a = W_NDimArray([num_items], dtype=dtype)
+    ai = a.create_iter()
+    for val in items:
+        a.dtype.setitem(a, ai.offset, val)
+        ai = ai.next(1)
     
     return space.wrap(a)
 
@@ -71,10 +73,12 @@
         raise OperationError(space.w_ValueError, space.wrap(
             "string is smaller than requested size"))
         
-    a = W_NDimArray(count, [count], dtype=dtype)
+    a = W_NDimArray([count], dtype=dtype)
+    ai = a.create_iter()
     for i in range(count):
         val = dtype.itemtype.runpack_str(s[i*itemsize:i*itemsize + itemsize])
-        a.dtype.setitem(a, i, val)
+        a.dtype.setitem(a, ai.offset, val)
+        ai = ai.next(1)
         
     return space.wrap(a)
 
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
@@ -156,7 +156,7 @@
             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)
+        result = W_NDimArray(shape, dtype)
         arr = AxisReduce(self.func, self.name, self.identity, obj.shape, dtype,
                          result, obj, dim)
         loop.compute(arr)
diff --git a/pypy/module/micronumpy/strides.py b/pypy/module/micronumpy/strides.py
--- a/pypy/module/micronumpy/strides.py
+++ b/pypy/module/micronumpy/strides.py
@@ -48,20 +48,20 @@
 def find_shape_and_elems(space, w_iterable, dtype):
     shape = [space.len_w(w_iterable)]
     batch = space.listview(w_iterable)
-    is_rec_type = dtype.is_record_type()
+    is_rec_type = dtype is not None and dtype.is_record_type()
     while True:
         new_batch = []
         if not batch:
             return shape, []
         if is_single_elem(space, batch[0], is_rec_type):
             for w_elem in batch:
-                if is_single_elem(space, w_elem, is_rec_type):
+                if not is_single_elem(space, w_elem, is_rec_type):
                     raise OperationError(space.w_ValueError, space.wrap(
                         "setting an array element with a sequence"))
             return shape, batch
         size = space.len_w(batch[0])
         for w_elem in batch:
-            if (not is_single_elem(space, w_elem, is_rec_type) or
+            if (is_single_elem(space, w_elem, is_rec_type) or
                 space.len_w(w_elem) != size):
                 raise OperationError(space.w_ValueError, space.wrap(
                     "setting an array element with a sequence"))
diff --git a/pypy/module/micronumpy/test/test_base.py b/pypy/module/micronumpy/test/test_base.py
--- a/pypy/module/micronumpy/test/test_base.py
+++ b/pypy/module/micronumpy/test/test_base.py
@@ -21,8 +21,8 @@
         float64_dtype = get_dtype_cache(space).w_float64dtype
         bool_dtype = get_dtype_cache(space).w_booldtype
 
-        ar = W_NDimArray(10, [10], dtype=float64_dtype)
-        ar2 = W_NDimArray(10, [10], dtype=float64_dtype)
+        ar = W_NDimArray([10], dtype=float64_dtype)
+        ar2 = W_NDimArray([10], dtype=float64_dtype)
         v1 = ar.descr_add(space, ar)
         v2 = ar.descr_add(space, Scalar(float64_dtype, W_Float64Box(2.0)))
         sig1 = v1.find_sig()
@@ -40,7 +40,7 @@
         v4 = ar.descr_add(space, ar)
         assert v1.find_sig() is v4.find_sig()
 
-        bool_ar = W_NDimArray(10, [10], dtype=bool_dtype)
+        bool_ar = W_NDimArray([10], dtype=bool_dtype)
         v5 = ar.descr_add(space, bool_ar)
         assert v5.find_sig() is not v1.find_sig()
         assert v5.find_sig() is not v2.find_sig()
@@ -57,7 +57,7 @@
     def test_slice_signature(self, space):
         float64_dtype = get_dtype_cache(space).w_float64dtype
 
-        ar = W_NDimArray(10, [10], dtype=float64_dtype)
+        ar = W_NDimArray([10], dtype=float64_dtype)
         v1 = ar.descr_getitem(space, space.wrap(slice(1, 3, 1)))
         v2 = ar.descr_getitem(space, space.wrap(slice(4, 6, 1)))
         assert v1.find_sig() is v2.find_sig()
diff --git a/pypy/module/micronumpy/test/test_iter.py b/pypy/module/micronumpy/test/test_iter.py
--- a/pypy/module/micronumpy/test/test_iter.py
+++ b/pypy/module/micronumpy/test/test_iter.py
@@ -49,17 +49,17 @@
         backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
         assert backstrides == [10, 4]
         i = ViewIterator(start, strides, backstrides, shape)
-        i = i.next_skip_x(2,2)
-        i = i.next_skip_x(2,2)
-        i = i.next_skip_x(2,2)
+        i = i.next_skip_ofs(2,2)
+        i = i.next_skip_ofs(2,2)
+        i = i.next_skip_ofs(2,2)
         assert i.offset == 6
         assert not i.done()
         assert i.indices == [1,1]
         #And for some big skips
-        i = i.next_skip_x(2,5)
+        i = i.next_skip_ofs(2,5)
         assert i.offset == 11
         assert i.indices == [2,1]
-        i = i.next_skip_x(2,5)
+        i = i.next_skip_ofs(2,5)
         # Note: the offset does not overflow but recycles,
         # this is good for broadcast
         assert i.offset == 1
@@ -72,17 +72,17 @@
         backstrides = [x * (y - 1) for x,y in zip(strides, shape)]
         assert backstrides == [2, 12]
         i = ViewIterator(start, strides, backstrides, shape)
-        i = i.next_skip_x(2,2)
-        i = i.next_skip_x(2,2)
-        i = i.next_skip_x(2,2)
+        i = i.next_skip_ofs(2,2)
+        i = i.next_skip_ofs(2,2)
+        i = i.next_skip_ofs(2,2)
         assert i.offset == 4
         assert i.indices == [1,1]
         assert not i.done()
-        i = i.next_skip_x(2,5)
+        i = i.next_skip_ofs(2,5)
         assert i.offset == 5
         assert i.indices == [2,1]
         assert not i.done()
-        i = i.next_skip_x(2,5)
+        i = i.next_skip_ofs(2,5)
         assert i.indices == [0,1]
         assert i.offset == 3
         assert i.done()
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
@@ -16,6 +16,9 @@
         def malloc(size):
             return None
 
+    def get_size(self):
+        return 1
+
 
 class TestNumArrayDirect(object):
     def newslice(self, *args):
@@ -31,17 +34,17 @@
         return self.space.newtuple(args_w)
 
     def test_strides_f(self):
-        a = W_NDimArray(100, [10, 5, 3], MockDtype(), 'F')
+        a = W_NDimArray([10, 5, 3], MockDtype(), 'F')
         assert a.strides == [1, 10, 50]
         assert a.backstrides == [9, 40, 100]
 
     def test_strides_c(self):
-        a = W_NDimArray(100, [10, 5, 3], MockDtype(), 'C')
+        a = W_NDimArray([10, 5, 3], MockDtype(), 'C')
         assert a.strides == [15, 3, 1]
         assert a.backstrides == [135, 12, 2]
 
     def test_create_slice_f(self):
-        a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'F')
+        a = W_NDimArray([10, 5, 3], MockDtype(), 'F')
         s = a.create_slice([Chunk(3, 0, 0, 1)])
         assert s.start == 3
         assert s.strides == [10, 50]
@@ -59,7 +62,7 @@
         assert s.shape == [10, 3]
 
     def test_create_slice_c(self):
-        a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'C')
+        a = W_NDimArray([10, 5, 3], MockDtype(), 'C')
         s = a.create_slice([Chunk(3, 0, 0, 1)])
         assert s.start == 45
         assert s.strides == [3, 1]
@@ -79,7 +82,7 @@
         assert s.shape == [10, 3]
 
     def test_slice_of_slice_f(self):
-        a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'F')
+        a = W_NDimArray([10, 5, 3], MockDtype(), 'F')
         s = a.create_slice([Chunk(5, 0, 0, 1)])
         assert s.start == 5
         s2 = s.create_slice([Chunk(3, 0, 0, 1)])
@@ -96,7 +99,7 @@
         assert s2.start == 1 * 15 + 2 * 3
 
     def test_slice_of_slice_c(self):
-        a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), order='C')
+        a = W_NDimArray([10, 5, 3], MockDtype(), order='C')
         s = a.create_slice([Chunk(5, 0, 0, 1)])
         assert s.start == 15 * 5
         s2 = s.create_slice([Chunk(3, 0, 0, 1)])
@@ -113,21 +116,21 @@
         assert s2.start == 1 * 15 + 2 * 3
 
     def test_negative_step_f(self):
-        a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'F')
+        a = W_NDimArray([10, 5, 3], MockDtype(), 'F')
         s = a.create_slice([Chunk(9, -1, -2, 5)])
         assert s.start == 9
         assert s.strides == [-2, 10, 50]
         assert s.backstrides == [-8, 40, 100]
 
     def test_negative_step_c(self):
-        a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), order='C')
+        a = W_NDimArray([10, 5, 3], MockDtype(), order='C')
         s = a.create_slice([Chunk(9, -1, -2, 5)])
         assert s.start == 135
         assert s.strides == [-30, 3, 1]
         assert s.backstrides == [-120, 12, 2]
 
     def test_index_of_single_item_f(self):
-        a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'F')
+        a = W_NDimArray([10, 5, 3], MockDtype(), 'F')
         r = a._index_of_single_item(self.space, self.newtuple(1, 2, 2))
         assert r == 1 + 2 * 10 + 2 * 50
         s = a.create_slice([Chunk(0, 10, 1, 10), Chunk(2, 0, 0, 1)])
@@ -137,7 +140,7 @@
         assert r == a._index_of_single_item(self.space, self.newtuple(1, 2, 1))
 
     def test_index_of_single_item_c(self):
-        a = W_NDimArray(10 * 5 * 3, [10, 5, 3], MockDtype(), 'C')
+        a = W_NDimArray([10, 5, 3], MockDtype(), 'C')
         r = a._index_of_single_item(self.space, self.newtuple(1, 2, 2))
         assert r == 1 * 3 * 5 + 2 * 3 + 2
         s = a.create_slice([Chunk(0, 10, 1, 10), Chunk(2, 0, 0, 1)])
@@ -1468,6 +1471,7 @@
         a = arange(12).reshape(3,4)
         b = a.T.flat
         b[6::2] = [-1, -2]
+        print a == [[0, 1, -1, 3], [4, 5, 6, -1], [8, 9, -2, 11]]
         assert (a == [[0, 1, -1, 3], [4, 5, 6, -1], [8, 9, -2, 11]]).all()
         b[0:2] = [[[100]]]
         assert(a[0,0] == 100)
diff --git a/pypy/module/micronumpy/test/test_zjit.py b/pypy/module/micronumpy/test/test_zjit.py
--- a/pypy/module/micronumpy/test/test_zjit.py
+++ b/pypy/module/micronumpy/test/test_zjit.py
@@ -503,7 +503,7 @@
                 dtype = float64_dtype
             else:
                 dtype = int32_dtype
-            ar = W_NDimArray(n, [n], dtype=dtype)
+            ar = W_NDimArray([n], dtype=dtype)
             i = 0
             while i < n:
                 ar.get_concrete().setitem(i, int32_dtype.box(7))
diff --git a/pypy/module/micronumpy/types.py b/pypy/module/micronumpy/types.py
--- a/pypy/module/micronumpy/types.py
+++ b/pypy/module/micronumpy/types.py
@@ -69,10 +69,9 @@
     #     exp = sin = cos = tan = arcsin = arccos = arctan = arcsinh = \
     #     arctanh = _unimplemented_ufunc
 
-    def malloc(self, length):
+    def malloc(self, size):
         # XXX find out why test_zjit explodes with tracking of allocations
-        return lltype.malloc(VOID_STORAGE,
-                             self.get_element_size() * length,
+        return lltype.malloc(VOID_STORAGE, size,
                              zero=True, flavor="raw",
                              track_allocation=False, add_memory_pressure=True)
 
@@ -140,8 +139,8 @@
 
     def fill(self, storage, width, box, start, stop, offset):
         value = self.unbox(box)
-        for i in xrange(start, stop):
-            self._write(storage, width, i, offset, value)
+        for i in xrange(start, stop, width):
+            self._write(storage, 1, i, offset, value)
 
     def runpack_str(self, s):
         return self.box(runpack(self.format_code, s))
@@ -667,7 +666,7 @@
         items_w = space.fixedview(w_item)
         # XXX optimize it out one day, but for now we just allocate an
         #     array
-        arr = W_NDimArray(1, [1], dtype)
+        arr = W_NDimArray([1], dtype)
         for i in range(len(items_w)):
             subdtype = dtype.fields[dtype.fieldnames[i]][1]
             ofs, itemtype = self.offsets_and_fields[i]


More information about the pypy-commit mailing list