def diagonal_array(space, arr, out, offset, axis1, axis2, shape):
out_iter, out_state = out.create_iter()
iter = PureShapeIter(shape, [])
shapelen_minus_1 = len(shape) - 1
assert shapelen_minus_1 >= 0
if axis1 < axis2:
a = axis1
b = axis2 - 1
else:
a = axis2
b = axis1 - 1
assert a >= 0
assert b >= 0
while not iter.done():
last_index = iter.indexes[-1]
if axis1 < axis2:
indexes = (iter.indexes[:a] + [last_index] +
iter.indexes[a:b] + [last_index + offset] +
iter.indexes[b:shapelen_minus_1])
else:
indexes = (iter.indexes[:a] + [last_index + offset] +
iter.indexes[a:b] + [last_index] +
iter.indexes[b:shapelen_minus_1])
out_iter.setitem(out_state, arr.getitem_index(space, indexes))
iter.next()
out_state = out_iter.next(out_state)
def setitem_array_int(space, arr, iter_shape, indexes_w, val_arr,
prefix_w):
shapelen = len(iter_shape)
indexlen = len(indexes_w)
prefixlen = len(prefix_w)
dtype = arr.get_dtype()
iter = PureShapeIter(iter_shape, indexes_w)
while not iter.done():
setitem_int_driver.jit_merge_point(shapelen=shapelen, indexlen=indexlen,
dtype=dtype, prefixlen=prefixlen)
# prepare the index
index_w = [None] * indexlen
for i in range(indexlen):
if iter.idx_w_i[i] is not None:
index_w[i] = iter.idx_w_i[i].getitem(iter.idx_w_s[i])
else:
index_w[i] = indexes_w[i]
w_idx = space.newtuple(prefix_w[:prefixlen] + iter.get_index(space,
shapelen))
if val_arr.is_scalar():
w_value = val_arr.get_scalar_value()
else:
w_value = val_arr.descr_getitem(space, w_idx)
arr.descr_setitem(space, space.newtuple(index_w), w_value)
iter.next()
def diagonal_array(space, arr, out, offset, axis1, axis2, shape):
out_iter, out_state = out.create_iter()
iter = PureShapeIter(shape, [])
shapelen_minus_1 = len(shape) - 1
assert shapelen_minus_1 >= 0
if axis1 < axis2:
a = axis1
b = axis2 - 1
else:
a = axis2
b = axis1 - 1
assert a >= 0
assert b >= 0
while not iter.done():
last_index = iter.indexes[-1]
if axis1 < axis2:
indexes = (iter.indexes[:a] + [last_index] + iter.indexes[a:b] +
[last_index + offset] +
iter.indexes[b:shapelen_minus_1])
else:
indexes = (iter.indexes[:a] + [last_index + offset] +
iter.indexes[a:b] + [last_index] +
iter.indexes[b:shapelen_minus_1])
out_iter.setitem(out_state, arr.getitem_index(space, indexes))
iter.next()
out_state = out_iter.next(out_state)
def getitem_array_int(space, arr, res, iter_shape, indexes_w, prefix_w):
shapelen = len(iter_shape)
prefixlen = len(prefix_w)
indexlen = len(indexes_w)
dtype = arr.get_dtype()
iter = PureShapeIter(iter_shape, indexes_w)
while not iter.done():
getitem_int_driver.jit_merge_point(shapelen=shapelen, indexlen=indexlen,
dtype=dtype, prefixlen=prefixlen)
# prepare the index
index_w = [None] * indexlen
for i in range(indexlen):
if iter.idx_w_i[i] is not None:
index_w[i] = iter.idx_w_i[i].getitem(iter.idx_w_s[i])
else:
index_w[i] = indexes_w[i]
res.descr_setitem(space, space.newtuple(prefix_w[:prefixlen] +
iter.get_index(space, shapelen)),
arr.descr_getitem(space, space.newtuple(index_w)))
iter.next()
return res
def setitem_array_int(space, arr, iter_shape, indexes_w, val_arr, prefix_w):
shapelen = len(iter_shape)
indexlen = len(indexes_w)
prefixlen = len(prefix_w)
dtype = arr.get_dtype()
iter = PureShapeIter(iter_shape, indexes_w)
while not iter.done():
setitem_int_driver.jit_merge_point(shapelen=shapelen,
indexlen=indexlen,
dtype=dtype,
prefixlen=prefixlen)
# prepare the index
index_w = [None] * indexlen
for i in range(indexlen):
if iter.idx_w_i[i] is not None:
index_w[i] = iter.idx_w_i[i].getitem(iter.idx_w_s[i])
else:
index_w[i] = indexes_w[i]
w_idx = space.newtuple(prefix_w[:prefixlen] +
iter.get_index(space, shapelen))
if val_arr.is_scalar():
w_value = val_arr.get_scalar_value()
else:
w_value = val_arr.descr_getitem(space, w_idx)
arr.descr_setitem(space, space.newtuple(index_w), w_value)
iter.next()
def getitem_array_int(space, arr, res, iter_shape, indexes_w, prefix_w):
shapelen = len(iter_shape)
prefixlen = len(prefix_w)
indexlen = len(indexes_w)
dtype = arr.get_dtype()
iter = PureShapeIter(iter_shape, indexes_w)
indexlen = len(indexes_w)
while not iter.done():
getitem_int_driver.jit_merge_point(shapelen=shapelen,
indexlen=indexlen,
dtype=dtype,
prefixlen=prefixlen)
# prepare the index
index_w = [None] * indexlen
for i in range(indexlen):
if iter.idx_w_i[i] is not None:
index_w[i] = iter.idx_w_i[i].getitem(iter.idx_w_s[i])
else:
index_w[i] = indexes_w[i]
res.descr_setitem(
space,
space.newtuple(prefix_w[:prefixlen] +
iter.get_index(space, shapelen)),
arr.descr_getitem(space, space.newtuple(index_w)))
iter.next()
return res
