Author: Ronan Lamy <ronan.l...@gmail.com>
Branch: indexing
Changeset: r78588:5663f55d1d87
Date: 2015-07-17 20:02 +0100
http://bitbucket.org/pypy/pypy/changeset/5663f55d1d87/
Log: implement ellipsis indexing
diff --git a/pypy/module/micronumpy/concrete.py
b/pypy/module/micronumpy/concrete.py
--- a/pypy/module/micronumpy/concrete.py
+++ b/pypy/module/micronumpy/concrete.py
@@ -222,18 +222,18 @@
if space.isinstance_w(w_idx, space.w_slice):
if len(self.get_shape()) == 0:
raise oefmt(space.w_ValueError, "cannot slice a 0-d array")
- return [SliceChunk(w_idx)]
+ return [SliceChunk(w_idx), EllipsisChunk()]
elif space.isinstance_w(w_idx, space.w_int):
- return [IntegerChunk(w_idx)]
+ return [IntegerChunk(w_idx), EllipsisChunk()]
elif isinstance(w_idx, W_NDimArray) and w_idx.is_scalar():
w_idx = w_idx.get_scalar_value().item(space)
if not space.isinstance_w(w_idx, space.w_int) and \
not space.isinstance_w(w_idx, space.w_bool):
raise OperationError(space.w_IndexError, space.wrap(
"arrays used as indices must be of integer (or boolean)
type"))
- return [IntegerChunk(w_idx)]
+ return [IntegerChunk(w_idx), EllipsisChunk()]
elif space.is_w(w_idx, space.w_None):
- return [NewAxisChunk()]
+ return [NewAxisChunk(), EllipsisChunk()]
result = []
i = 0
has_ellipsis = False
@@ -254,6 +254,8 @@
else:
result.append(IntegerChunk(w_item))
i += 1
+ if not has_ellipsis:
+ result.append(EllipsisChunk())
return result
def descr_getitem(self, space, orig_arr, w_index):
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
@@ -59,12 +59,7 @@
backstride = base_stride * max(0, length - 1) * step
return start, length, stride, backstride
-
class NewAxisChunk(Chunk):
- start = 0
- stop = 1
- step = 1
- lgt = 1
input_dim = 0
out_dim = 1
@@ -75,9 +70,15 @@
return 0, 1, 0, 0
class EllipsisChunk(BaseChunk):
+ input_dim = 0
+ out_dim = 0
def __init__(self):
pass
+ def compute(self, space, base_length, base_stride):
+ backstride = base_stride * max(0, base_length - 1)
+ return 0, base_length, base_stride, backstride
+
def new_view(space, w_arr, chunks):
arr = w_arr.implementation
@@ -127,38 +128,48 @@
@jit.look_inside_iff(lambda space, shape, start, strides, backstrides, chunks:
jit.isconstant(len(chunks)))
def calculate_slice_strides(space, shape, start, strides, backstrides, chunks):
+ """
+ Note: `chunks` must contain exactly one EllipsisChunk object.
+ """
size = 0
+ used_dims = 0
for chunk in chunks:
+ used_dims += chunk.input_dim
size += chunk.out_dim
- rstrides = [0] * size
- rbackstrides = [0] * size
+ if used_dims > len(shape):
+ raise oefmt(space.w_IndexError, "too many indices for array")
+ else:
+ extra_dims = len(shape) - used_dims
+ rstrides = [0] * (size + extra_dims)
+ rbackstrides = [0] * (size + extra_dims)
rstart = start
- rshape = [0] * size
- i = -1 # index of the current dimension in the input array
+ rshape = [0] * (size + extra_dims)
+ rstart = start
+ i = 0 # index of the current dimension in the input array
j = 0 # index of the current dimension in the result view
for chunk in chunks:
- i += chunk.input_dim
if isinstance(chunk, NewAxisChunk):
rshape[j] = 1
j += 1
continue
- try:
- s_i = strides[i]
- except IndexError:
+ elif isinstance(chunk, EllipsisChunk):
+ for k in range(extra_dims):
+ start, length, stride, backstride = chunk.compute(
+ space, shape[i], strides[i])
+ rshape[j] = length
+ rstrides[j] = stride
+ rbackstrides[j] = backstride
+ j += 1
+ i += 1
continue
start, length, stride, backstride = chunk.compute(space, shape[i],
strides[i])
if chunk.out_dim == 1:
+ rshape[j] = length
rstrides[j] = stride
rbackstrides[j] = backstride
- rshape[j] = length
j += chunk.out_dim
- rstart += s_i * start
- # add a reminder
- s = i + 1
- assert s >= 0
- rstrides += strides[s:]
- rbackstrides += backstrides[s:]
- rshape += shape[s:]
+ rstart += strides[i] * start
+ i += chunk.input_dim
return rshape, rstart, rstrides, rbackstrides
diff --git a/pypy/module/micronumpy/test/test_ndarray.py
b/pypy/module/micronumpy/test/test_ndarray.py
--- a/pypy/module/micronumpy/test/test_ndarray.py
+++ b/pypy/module/micronumpy/test/test_ndarray.py
@@ -4,7 +4,7 @@
from pypy.conftest import option
from pypy.module.micronumpy.appbridge import get_appbridge_cache
-from pypy.module.micronumpy.strides import Chunk, new_view
+from pypy.module.micronumpy.strides import Chunk, new_view, EllipsisChunk
from pypy.module.micronumpy.ndarray import W_NDimArray
from pypy.module.micronumpy.test.test_base import BaseNumpyAppTest
@@ -22,6 +22,8 @@
def create_slice(space, a, chunks):
+ if not any(isinstance(c, EllipsisChunk) for c in chunks):
+ chunks.append(EllipsisChunk())
return new_view(space, W_NDimArray(a), chunks).implementation
@@ -2490,6 +2492,9 @@
a = np.arange(6).reshape(2, 3)
if '__pypy__' in sys.builtin_module_names:
raises(ValueError, "a[..., ...]")
+ b = a [..., 0]
+ assert (b == [0, 3]).all()
+ assert b.base is a
def test_empty_indexing(self):
import numpy as np
_______________________________________________
pypy-commit mailing list
pypy-commit@python.org
https://mail.python.org/mailman/listinfo/pypy-commit
