def sort(arr, space, w_axis):
if w_axis is space.w_None:
# note that it's fine to pass None here as we're not going
# to pass the result around (None is the link to base in slices)
arr = arr.reshape(None, [arr.get_size()])
axis = 0
elif w_axis is None:
axis = -1
else:
axis = space.int_w(w_axis)
with arr as storage:
if len(arr.get_shape()) == 1:
r = Repr(arr.strides[0], arr.get_size(), storage,
arr.start)
ArgSort(r).sort()
else:
shape = arr.get_shape()
if axis < 0:
axis = len(shape) + axis
if axis < 0 or axis >= len(shape):
raise oefmt(space.w_IndexError, "Wrong axis %d", axis)
arr_iter = AllButAxisIter(arr, axis)
arr_state = arr_iter.reset()
stride_size = arr.strides[axis]
axis_size = arr.shape[axis]
while not arr_iter.done(arr_state):
r = Repr(stride_size, axis_size, storage, arr_state.offset)
ArgSort(r).sort()
arr_state = arr_iter.next(arr_state)
def sort(arr, space, w_axis):
if w_axis is space.w_None:
# note that it's fine to pass None here as we're not going
# to pass the result around (None is the link to base in slices)
arr = arr.reshape(None, [arr.get_size()])
axis = 0
elif w_axis is None:
axis = -1
else:
axis = space.int_w(w_axis)
with arr as storage:
if len(arr.get_shape()) == 1:
r = Repr(arr.strides[0], arr.get_size(), storage, arr.start)
ArgSort(r).sort()
else:
shape = arr.get_shape()
if axis < 0:
axis = len(shape) + axis
if axis < 0 or axis >= len(shape):
raise oefmt(space.w_IndexError, "Wrong axis %d", axis)
arr_iter = AllButAxisIter(arr, axis)
arr_state = arr_iter.reset()
stride_size = arr.strides[axis]
axis_size = arr.shape[axis]
while not arr_iter.done(arr_state):
r = Repr(stride_size, axis_size, storage, arr_state.offset)
ArgSort(r).sort()
arr_state = arr_iter.next(arr_state)
def multidim_dot(space, left, right, result, dtype, right_critical_dim):
''' assumes left, right are concrete arrays
given left.shape == [3, 5, 7],
right.shape == [2, 7, 4]
then
result.shape == [3, 5, 2, 4]
broadcast shape should be [3, 5, 2, 7, 4]
result should skip dims 3 which is len(result_shape) - 1
(note that if right is 1d, result should
skip len(result_shape))
left should skip 2, 4 which is a.ndims-1 + range(right.ndims)
except where it==(right.ndims-2)
right should skip 0, 1
'''
left_shape = left.get_shape()
right_shape = right.get_shape()
left_impl = left.implementation
right_impl = right.implementation
assert left_shape[-1] == right_shape[right_critical_dim]
assert result.get_dtype() == dtype
outi, outs = result.create_iter()
outi.track_index = False
lefti = AllButAxisIter(left_impl, len(left_shape) - 1)
righti = AllButAxisIter(right_impl, right_critical_dim)
lefts = lefti.reset()
rights = righti.reset()
n = left_impl.shape[-1]
s1 = left_impl.strides[-1]
s2 = right_impl.strides[right_critical_dim]
while not lefti.done(lefts):
while not righti.done(rights):
oval = outi.getitem(outs)
i1 = lefts.offset
i2 = rights.offset
i = 0
while i < n:
i += 1
dot_driver.jit_merge_point(dtype=dtype)
lval = left_impl.getitem(i1).convert_to(space, dtype)
rval = right_impl.getitem(i2).convert_to(space, dtype)
oval = dtype.itemtype.add(oval, dtype.itemtype.mul(lval, rval))
i1 += s1
i2 += s2
outi.setitem(outs, oval)
outs = outi.next(outs)
rights = righti.next(rights)
rights = righti.reset(rights)
lefts = lefti.next(lefts)
return result
def argsort(arr, space, w_axis):
if w_axis is space.w_None:
# note that it's fine ot pass None here as we're not going
# to pass the result around (None is the link to base in slices)
if arr.get_size() > 0:
arr = arr.reshape(None, [arr.get_size()])
axis = 0
elif w_axis is None:
axis = -1
else:
axis = space.int_w(w_axis)
# create array of indexes
dtype = descriptor.get_dtype_cache(space).w_longdtype
index_arr = W_NDimArray.from_shape(space, arr.get_shape(), dtype)
with index_arr.implementation as storage, arr as arr_storage:
if len(arr.get_shape()) == 1:
for i in range(arr.get_size()):
raw_storage_setitem(storage, i * INT_SIZE, i)
r = Repr(INT_SIZE, arr.strides[0], arr.get_size(), arr_storage,
storage, 0, arr.start)
ArgSort(r).sort()
else:
shape = arr.get_shape()
if axis < 0:
axis = len(shape) + axis
if axis < 0 or axis >= len(shape):
raise oefmt(space.w_IndexError, "Wrong axis %d", axis)
arr_iter = AllButAxisIter(arr, axis)
arr_state = arr_iter.reset()
index_impl = index_arr.implementation
index_iter = AllButAxisIter(index_impl, axis)
index_state = index_iter.reset()
stride_size = arr.strides[axis]
index_stride_size = index_impl.strides[axis]
axis_size = arr.shape[axis]
while not arr_iter.done(arr_state):
for i in range(axis_size):
raw_storage_setitem(
storage,
i * index_stride_size + index_state.offset, i)
r = Repr(index_stride_size, stride_size, axis_size,
arr_storage, storage, index_state.offset,
arr_state.offset)
ArgSort(r).sort()
arr_state = arr_iter.next(arr_state)
index_state = index_iter.next(index_state)
return index_arr
def argsort(arr, space, w_axis):
if w_axis is space.w_None:
# note that it's fine ot pass None here as we're not going
# to pass the result around (None is the link to base in slices)
if arr.get_size() > 0:
arr = arr.reshape(None, [arr.get_size()])
axis = 0
elif w_axis is None:
axis = -1
else:
axis = space.int_w(w_axis)
# create array of indexes
dtype = descriptor.get_dtype_cache(space).w_longdtype
index_arr = W_NDimArray.from_shape(space, arr.get_shape(), dtype)
with index_arr.implementation as storage, arr as arr_storage:
if len(arr.get_shape()) == 1:
for i in range(arr.get_size()):
raw_storage_setitem(storage, i * INT_SIZE, i)
r = Repr(INT_SIZE, arr.strides[0], arr.get_size(), arr_storage,
storage, 0, arr.start)
ArgSort(r).sort()
else:
shape = arr.get_shape()
if axis < 0:
axis = len(shape) + axis
if axis < 0 or axis >= len(shape):
raise oefmt(space.w_IndexError, "Wrong axis %d", axis)
arr_iter = AllButAxisIter(arr, axis)
arr_state = arr_iter.reset()
index_impl = index_arr.implementation
index_iter = AllButAxisIter(index_impl, axis)
index_state = index_iter.reset()
stride_size = arr.strides[axis]
index_stride_size = index_impl.strides[axis]
axis_size = arr.shape[axis]
while not arr_iter.done(arr_state):
for i in range(axis_size):
raw_storage_setitem(storage, i * index_stride_size +
index_state.offset, i)
r = Repr(index_stride_size, stride_size, axis_size,
arr_storage, storage, index_state.offset, arr_state.offset)
ArgSort(r).sort()
arr_state = arr_iter.next(arr_state)
index_state = index_iter.next(index_state)
return index_arr
def multidim_dot(space, left, right, result, dtype, right_critical_dim):
''' assumes left, right are concrete arrays
given left.shape == [3, 5, 7],
right.shape == [2, 7, 4]
then
result.shape == [3, 5, 2, 4]
broadcast shape should be [3, 5, 2, 7, 4]
result should skip dims 3 which is len(result_shape) - 1
(note that if right is 1d, result should
skip len(result_shape))
left should skip 2, 4 which is a.ndims-1 + range(right.ndims)
except where it==(right.ndims-2)
right should skip 0, 1
'''
left_shape = left.get_shape()
right_shape = right.get_shape()
left_impl = left.implementation
right_impl = right.implementation
assert left_shape[-1] == right_shape[right_critical_dim]
assert result.get_dtype() == dtype
outi, outs = result.create_iter()
outi.track_index = False
lefti = AllButAxisIter(left_impl, len(left_shape) - 1)
righti = AllButAxisIter(right_impl, right_critical_dim)
lefts = lefti.reset()
rights = righti.reset()
n = left_impl.shape[-1]
s1 = left_impl.strides[-1]
s2 = right_impl.strides[right_critical_dim]
while not lefti.done(lefts):
while not righti.done(rights):
oval = outi.getitem(outs)
i1 = lefts.offset
i2 = rights.offset
i = 0
while i < n:
i += 1
dot_driver.jit_merge_point(dtype=dtype)
lval = left_impl.getitem(i1).convert_to(space, dtype)
rval = right_impl.getitem(i2).convert_to(space, dtype)
oval = dtype.itemtype.add(oval, dtype.itemtype.mul(lval, rval))
i1 += s1
i2 += s2
outi.setitem(outs, oval)
outs = outi.next(outs)
rights = righti.next(rights)
rights = righti.reset(rights)
lefts = lefti.next(lefts)
return result
