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 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 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