diff --git a/numba/core/typing/arraydecl.py b/numba/core/typing/arraydecl.py index 60d4d0201c4..128aae5aca2 100644 --- a/numba/core/typing/arraydecl.py +++ b/numba/core/typing/arraydecl.py @@ -33,46 +33,55 @@ def get_array_index_type(ary, idx): right_indices = [] ellipsis_met = False advanced = False - has_integer = False if not isinstance(idx, types.BaseTuple): idx = [idx] + in_subspace = False + num_subspaces = 0 + array_indices = [] + # Walk indices for ty in idx: if ty is types.ellipsis: if ellipsis_met: - raise NumbaTypeError("only one ellipsis allowed in array index " + raise NumbaTypeError("Only one ellipsis allowed in array indices " "(got %s)" % (idx,)) ellipsis_met = True + in_subspace = False elif isinstance(ty, types.SliceType): - pass + in_subspace = False elif isinstance(ty, types.Integer): # Normalize integer index ty = types.intp if ty.signed else types.uintp # Integer indexing removes the given dimension ndim -= 1 - has_integer = True + if not in_subspace: + num_subspaces += 1 + in_subspace = True elif (isinstance(ty, types.Array) and ty.ndim == 0 and isinstance(ty.dtype, types.Integer)): # 0-d array used as integer index ndim -= 1 - has_integer = True + if not in_subspace: + num_subspaces += 1 + in_subspace = True elif (isinstance(ty, types.Array) - and ty.ndim == 1 and isinstance(ty.dtype, (types.Integer, types.Boolean))): - if advanced or has_integer: - # We don't support the complicated combination of - # advanced indices (and integers are considered part - # of them by Numpy). - msg = "only one advanced index supported" - raise NumbaNotImplementedError(msg) + array_indices.append(ty.ndim) advanced = True + ndim -= 1 + if not in_subspace: + num_subspaces += 1 + in_subspace = True else: - raise NumbaTypeError("unsupported array index type %s in %s" + raise NumbaTypeError("Unsupported array index type %s in %s" % (ty, idx)) (right_indices if ellipsis_met else left_indices).append(ty) + if advanced: + ndim += max(array_indices) + # Only Numpy arrays support advanced indexing if advanced and not isinstance(ary, types.Array): return diff --git a/numba/np/arrayobj.py b/numba/np/arrayobj.py index 80080464fd9..4ccf48bb2da 100644 --- a/numba/np/arrayobj.py +++ b/numba/np/arrayobj.py @@ -634,12 +634,9 @@ def get_index_bounds(self): def loop_head(self): """ - Start indexation loop. Return a (index, count) tuple. + Start indexation loop. Returns a index. *index* is an integer LLVM value representing the index over this dimension. - *count* is either an integer LLVM value representing the current - iteration count, or None if this dimension should be omitted from - the indexation result. """ raise NotImplementedError @@ -690,7 +687,7 @@ def loop_head(self): with builder.if_then(builder.icmp_signed('>=', cur_index, self.size), likely=False): builder.branch(self.bb_end) - return cur_index, cur_index + return cur_index def loop_tail(self): builder = self.builder @@ -725,7 +722,7 @@ def get_index_bounds(self): return (self.idx, self.builder.add(self.idx, self.get_size())) def loop_head(self): - return self.idx, None + return self.idx def loop_tail(self): pass @@ -736,19 +733,48 @@ class IntegerArrayIndexer(Indexer): Compute indices from an array of integer indices. """ - def __init__(self, context, builder, idxty, idxary, size): + def __init__(self, context, builder, idxty, idxary, size, subspace_allocated, global_ary_idx): self.context = context self.builder = builder - self.idxty = idxty - self.idxary = idxary + self.subspace_allocated = subspace_allocated + self.global_ary_idx = global_ary_idx + + self.idx_shape = cgutils.unpack_tuple(builder, idxary.shape) self.size = size - assert idxty.ndim == 1 self.ll_intp = self.context.get_value_type(types.intp) + if idxty.ndim > 1: + def flat_imp_nocopy(ary): + return ary.reshape(ary.size) + + def flat_imp_copy(ary): + return ary.copy().reshape(ary.size) + + # If the index array is contigous, use the nocopy version + if idxty.is_contig: + flat_imp = flat_imp_nocopy + # otherwise, use copy version since we don't support + # reshaping non-contigous arrays + else: + flat_imp = flat_imp_copy + + retty = types.Array(idxty.dtype, 1, idxty.layout, readonly=True) + sig = signature(retty, idxty) + res = context.compile_internal(builder, flat_imp, sig, + (idxary._getvalue(),)) + self.idxty = retty + self.idxary = make_array(retty)(context, builder, res) + else: + self.idxty = idxty + self.idxary = idxary + + assert self.idxty.ndim == 1 + def prepare(self): builder = self.builder - self.idx_size = cgutils.unpack_tuple(builder, self.idxary.shape)[0] - self.idx_index = cgutils.alloca_once(builder, self.ll_intp) + self.idx_size = self.ll_intp(1) + for _shape in self.idx_shape: + self.idx_size = self.builder.mul(self.idx_size, _shape) self.bb_start = builder.append_basic_block() self.bb_end = builder.append_basic_block() @@ -756,7 +782,7 @@ def get_size(self): return self.idx_size def get_shape(self): - return (self.idx_size,) + return tuple(self.idx_shape) def get_index_bounds(self): # Pessimal heuristic, as we don't want to scan for the min and max @@ -765,15 +791,16 @@ def get_index_bounds(self): def loop_head(self): builder = self.builder # Initialize loop variable - self.builder.store(Constant(self.ll_intp, 0), self.idx_index) builder.branch(self.bb_start) builder.position_at_end(self.bb_start) - cur_index = builder.load(self.idx_index) - with builder.if_then( - builder.icmp_signed('>=', cur_index, self.idx_size), - likely=False - ): - builder.branch(self.bb_end) + cur_index = builder.load(self.global_ary_idx) + if not self.subspace_allocated: + with builder.if_then( + builder.icmp_signed('>=', cur_index, self.idx_size), + likely=False + ): + self.builder.store(Constant(self.ll_intp, 0), self.global_ary_idx) + builder.branch(self.bb_end) # Load the actual index from the array of indices index = _getitem_array_single_int( self.context, builder, self.idxty.dtype, self.idxty, self.idxary, @@ -781,14 +808,17 @@ def loop_head(self): ) index = fix_integer_index(self.context, builder, self.idxty.dtype, index, self.size) - return index, cur_index + return index def loop_tail(self): builder = self.builder - next_index = cgutils.increment_index(builder, - builder.load(self.idx_index)) - builder.store(next_index, self.idx_index) - builder.branch(self.bb_start) + if not self.subspace_allocated: + next_index = cgutils.increment_index(builder, + builder.load(self.global_ary_idx)) + builder.store(next_index, self.global_ary_idx) + builder.branch(self.bb_start) + else: + builder.branch(self.bb_end) builder.position_at_end(self.bb_end) @@ -810,7 +840,6 @@ def prepare(self): builder = self.builder self.size = cgutils.unpack_tuple(builder, self.idxary.shape)[0] self.idx_index = cgutils.alloca_once(builder, self.ll_intp) - self.count = cgutils.alloca_once(builder, self.ll_intp) self.bb_start = builder.append_basic_block() self.bb_tail = builder.append_basic_block() self.bb_end = builder.append_basic_block() @@ -842,11 +871,9 @@ def loop_head(self): builder = self.builder # Initialize loop variable self.builder.store(self.zero, self.idx_index) - self.builder.store(self.zero, self.count) builder.branch(self.bb_start) builder.position_at_end(self.bb_start) cur_index = builder.load(self.idx_index) - cur_count = builder.load(self.count) with builder.if_then(builder.icmp_signed('>=', cur_index, self.size), likely=False): builder.branch(self.bb_end) @@ -857,10 +884,7 @@ def loop_head(self): ) with builder.if_then(builder.not_(pred)): builder.branch(self.bb_tail) - # Increment the count for next iteration - next_count = cgutils.increment_index(builder, cur_count) - builder.store(next_count, self.count) - return cur_index, cur_count + return cur_index def loop_tail(self): builder = self.builder @@ -899,7 +923,6 @@ def prepare(self): self.is_step_negative = cgutils.is_neg_int(builder, self.slice.step) # Create loop entities self.index = cgutils.alloca_once(builder, self.ll_intp) - self.count = cgutils.alloca_once(builder, self.ll_intp) self.bb_start = builder.append_basic_block() self.bb_end = builder.append_basic_block() @@ -917,11 +940,9 @@ def loop_head(self): builder = self.builder # Initialize loop variable self.builder.store(self.slice.start, self.index) - self.builder.store(self.zero, self.count) builder.branch(self.bb_start) builder.position_at_end(self.bb_start) cur_index = builder.load(self.index) - cur_count = builder.load(self.count) is_finished = builder.select(self.is_step_negative, builder.icmp_signed('<=', cur_index, self.slice.stop), @@ -929,15 +950,13 @@ def loop_head(self): self.slice.stop)) with builder.if_then(is_finished, likely=False): builder.branch(self.bb_end) - return cur_index, cur_count + return cur_index def loop_tail(self): builder = self.builder next_index = builder.add(builder.load(self.index), self.slice.step, flags=['nsw']) builder.store(next_index, self.index) - next_count = cgutils.increment_index(builder, builder.load(self.count)) - builder.store(next_count, self.count) builder.branch(self.bb_start) builder.position_at_end(self.bb_end) @@ -947,17 +966,20 @@ class FancyIndexer(object): Perform fancy indexing on the given array. """ - def __init__(self, context, builder, aryty, ary, index_types, indices): + def __init__(self, context, builder, aryty, ary, index_types, indices, subspace_shape_tuple): self.context = context self.builder = builder self.aryty = aryty self.shapes = cgutils.unpack_tuple(builder, ary.shape, aryty.ndim) self.strides = cgutils.unpack_tuple(builder, ary.strides, aryty.ndim) self.ll_intp = self.context.get_value_type(types.intp) - + self.subspace_shape = subspace_shape_tuple + self.global_ary_idx = cgutils.alloca_once(builder, self.ll_intp) + self.builder.store(Constant(self.ll_intp, 0), self.global_ary_idx) indexers = [] ax = 0 + subspace_allocated=False for indexval, idxty in zip(indices, index_types): if idxty is types.ellipsis: # Fill up missing dimensions at the middle @@ -984,7 +1006,10 @@ def __init__(self, context, builder, aryty, ary, index_types, indices): if isinstance(idxty.dtype, types.Integer): indexer = IntegerArrayIndexer(context, builder, idxty, idxary, - self.shapes[ax]) + self.shapes[ax], + subspace_allocated, + self.global_ary_idx) + subspace_allocated = True elif isinstance(idxty.dtype, types.Boolean): indexer = BooleanArrayIndexer(context, builder, idxty, idxary) @@ -1009,7 +1034,18 @@ def prepare(self): for i in self.indexers: i.prepare() # Compute the resulting shape - self.indexers_shape = sum([i.get_shape() for i in self.indexers], ()) + res_shape = [] + subspace_added=False + for i in self.indexers: + # Shape of subspace i.e. cumulative shape of indexing + # arrays only needs to be considered once + if isinstance(i, IntegerArrayIndexer): + if not subspace_added: + res_shape.append(self.subspace_shape) + subspace_added = True + else: + res_shape.append(i.get_shape()) + self.indexers_shape = sum(res_shape, ()) def get_shape(self): """ @@ -1056,14 +1092,40 @@ def get_offset_bounds(self, strides, itemsize): return lower, upper def begin_loops(self): - indices, counts = zip(*(i.loop_head() for i in self.indexers)) - return indices, counts + indices = tuple(i.loop_head() for i in self.indexers) + return indices def end_loops(self): for i in reversed(self.indexers): i.loop_tail() +def get_bdcast_idx(context, builder, array_indices): + max_dims = max([ary[2].ndim for ary in array_indices]) + + def bdcast_idx_shapes(*args): + return np.broadcast_shapes(*args) + + inpty = types.StarArgTuple(tuple(types.UniTuple(types.intp, count=ary_idx[2].ndim) for ary_idx in array_indices)) + retty = types.UniTuple(types.intp, count=max_dims) + subspace_shape = context.compile_internal(builder, bdcast_idx_shapes, signature(retty, inpty), + (cgutils.pack_struct(builder, tuple([ary_idx[3].shape for ary_idx in array_indices])),)) + + bdcast_indices = [] + + def bdcast_array(ary, shape): + return np.broadcast_to(ary, shape) + + for i, idx, idxty, _ in array_indices: + inpty = (idxty, types.UniTuple(types.intp, count=max_dims)) + retty = types.Array(idxty.dtype, max_dims, 'A', readonly=True) + bdcast_idx = context.compile_internal(builder, bdcast_array, signature(retty, *inpty), + (idx, subspace_shape)) + bdcast_indices.append((i, bdcast_idx, retty)) + subspace_shape = tuple(cgutils.unpack_tuple(builder, subspace_shape)) + return bdcast_indices, subspace_shape + + def fancy_getitem(context, builder, sig, args, aryty, ary, index_types, indices): @@ -1071,8 +1133,24 @@ def fancy_getitem(context, builder, sig, args, strides = cgutils.unpack_tuple(builder, ary.strides) data = ary.data + array_indices = [] + for i, idxty in enumerate(index_types): + idx = indices[i] + if isinstance(idxty, types.Array): + idx_make = make_array(idxty)(context, builder, idx) + array_indices.append((i, idx, idxty, idx_make)) + + bdcast_indices, subspace_shape_tuple = \ + get_bdcast_idx(context, builder, array_indices) + + indices = list(indices) + index_types = list(index_types) + for i, bdcast_idx, bdcast_idxty in bdcast_indices: + indices[i] = bdcast_idx + index_types[i] = bdcast_idxty + indexer = FancyIndexer(context, builder, aryty, ary, - index_types, indices) + index_types, indices, subspace_shape_tuple) indexer.prepare() # Construct output array @@ -1085,7 +1163,7 @@ def fancy_getitem(context, builder, sig, args, context.get_constant(types.intp, 0)) # Loop on source and copy to destination - indices, _ = indexer.begin_loops() + indices = indexer.begin_loops() # No need to check for wraparound, as the indexers all ensure # a positive index is returned. @@ -1236,7 +1314,6 @@ def maybe_copy_source(context, builder, use_copy, builder.store(builder.load(src_ptr), dest_ptr) def src_getitem(source_indices): - assert len(source_indices) == srcty.ndim src_ptr = cgutils.alloca_once(builder, ptrty) with builder.if_else(use_copy, likely=False) as (if_copy, otherwise): with if_copy: @@ -1620,7 +1697,6 @@ def fancy_setslice(context, builder, sig, args, index_types, indices): srcty, src = _broadcast_to_shape(context, builder, srcty, src, index_shape) src_shapes = cgutils.unpack_tuple(builder, src.shape) - src_strides = cgutils.unpack_tuple(builder, src.strides) src_data = src.data # Check shapes are equal @@ -1635,24 +1711,6 @@ def fancy_setslice(context, builder, sig, args, index_types, indices): msg = "cannot assign slice from input of different size" context.call_conv.return_user_exc(builder, ValueError, (msg,)) - # Check for array overlap - src_start, src_end = get_array_memory_extents(context, builder, srcty, - src, src_shapes, - src_strides, src_data) - - dest_lower, dest_upper = indexer.get_offset_bounds(dest_strides, - ary.itemsize) - dest_start, dest_end = compute_memory_extents(context, builder, - dest_lower, dest_upper, - dest_data) - - use_copy = extents_may_overlap(context, builder, src_start, src_end, - dest_start, dest_end) - - src_getitem, src_cleanup = maybe_copy_source(context, builder, use_copy, - srcty, src, src_shapes, - src_strides, src_data) - elif isinstance(srcty, types.Sequence): src_dtype = srcty.dtype @@ -1667,51 +1725,54 @@ def fancy_setslice(context, builder, sig, args, index_types, indices): with builder.if_then(shape_error, likely=False): msg = "cannot assign slice from input of different size" context.call_conv.return_user_exc(builder, ValueError, (msg,)) - - def src_getitem(source_indices): - idx, = source_indices - getitem_impl = context.get_function( - operator.getitem, - signature(src_dtype, srcty, types.intp), - ) - return getitem_impl(builder, (src, idx)) - - def src_cleanup(): - pass - else: # Source is a scalar (broadcast or not, depending on destination # shape). src_dtype = srcty - def src_getitem(source_indices): - return src + def flat_imp_nocopy(ary): + return ary.reshape(ary.size) - def src_cleanup(): - pass + def flat_imp_copy(ary): + return ary.copy().reshape(ary.size) + # If the source array is contigous, use the nocopy version + if srcty.is_contig: + flat_imp = flat_imp_nocopy + # otherwise, use copy version since we don't support + # reshaping non-contigous arrays + else: + flat_imp = flat_imp_copy + + retty = types.Array(srcty.dtype, 1, srcty.layout, readonly=True) + sig = signature(retty, srcty) + src_flat_instr = context.compile_internal(builder, flat_imp, sig, + (src._getvalue(),)) + src_flat = make_array(retty)(context, builder, src_flat_instr) + src_data = src_flat.data + src_idx = cgutils.alloca_once_value(builder, + context.get_constant(types.intp, 0)) # Loop on destination and copy from source to destination - dest_indices, counts = indexer.begin_loops() - - # Source is iterated in natural order - source_indices = tuple(c for c in counts if c is not None) - val = src_getitem(source_indices) - - # Cast to the destination dtype (cross-dtype slice assignment is allowed) - val = context.cast(builder, val, src_dtype, aryty.dtype) + dest_indices = indexer.begin_loops() # No need to check for wraparound, as the indexers all ensure # a positive index is returned. dest_ptr = cgutils.get_item_pointer2(context, builder, dest_data, dest_shapes, dest_strides, aryty.layout, dest_indices, - wraparound=False, - boundscheck=context.enable_boundscheck) + wraparound=False) + + cur = builder.load(src_idx) + ptr = builder.gep(src_data, [cur]) + val = builder.load(ptr) + val = context.cast(builder, val, src_dtype, aryty.dtype) store_item(context, builder, aryty, val, dest_ptr) + next_idx = cgutils.increment_index(builder, cur) + builder.store(next_idx, src_idx) indexer.end_loops() - src_cleanup() + context.nrt.decref(builder, retty, src_flat_instr) return context.get_dummy_value() diff --git a/numba/tests/test_array_manipulation.py b/numba/tests/test_array_manipulation.py index 82687a6373a..7a4abf1f177 100644 --- a/numba/tests/test_array_manipulation.py +++ b/numba/tests/test_array_manipulation.py @@ -648,7 +648,7 @@ def test_add_axis2(self, flags=enable_pyobj_flags): def test_add_axis2_npm(self): with self.assertTypingError() as raises: self.test_add_axis2(flags=no_pyobj_flags) - self.assertIn("unsupported array index type none in", + self.assertIn("Unsupported array index type none in", str(raises.exception)) def test_bad_index_npm(self): @@ -658,13 +658,13 @@ def test_bad_index_npm(self): arraytype2 = types.Array(types.int32, 2, 'C') compile_isolated(bad_index, (arraytype1, arraytype2), flags=no_pyobj_flags) - self.assertIn('unsupported array index type', str(raises.exception)) + self.assertIn('Unsupported array index type', str(raises.exception)) def test_bad_float_index_npm(self): with self.assertTypingError() as raises: compile_isolated(bad_float_index, (types.Array(types.float64, 2, 'C'),)) - self.assertIn('unsupported array index type float64', + self.assertIn('Unsupported array index type float64', str(raises.exception)) def test_fill_diagonal_basic(self): diff --git a/numba/tests/test_fancy_indexing.py b/numba/tests/test_fancy_indexing.py index c9402a8e2e3..f0aeddefc73 100644 --- a/numba/tests/test_fancy_indexing.py +++ b/numba/tests/test_fancy_indexing.py @@ -3,183 +3,204 @@ import numpy as np import unittest -from numba import jit, typeof +from numba import njit, typeof from numba.core import types from numba.core.errors import TypingError from numba.tests.support import MemoryLeakMixin, TestCase, tag -def getitem_usecase(a, b): - return a[b] - -def setitem_usecase(a, idx, b): - a[idx] = b - -def np_take(A, indices): - return np.take(A, indices) - -def np_take_kws(A, indices, axis): - return np.take(A, indices, axis=axis) - class TestFancyIndexing(MemoryLeakMixin, TestCase): - - def generate_advanced_indices(self, N, many=True): - choices = [np.int16([0, N - 1, -2])] - if many: - choices += [np.uint16([0, 1, N - 1]), - np.bool_([0, 1, 1, 0])] - return choices - - def generate_basic_index_tuples(self, N, maxdim, many=True): - """ - Generate basic index tuples with 0 to *maxdim* items. - """ - # Note integers can be considered advanced indices in certain - # cases, so we avoid them here. - # See "Combining advanced and basic indexing" - # in http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html - if many: - choices = [slice(None, None, None), - slice(1, N - 1, None), - slice(0, None, 2), - slice(N - 1, None, -2), - slice(-N + 1, -1, None), - slice(-1, -N, -2), - ] - else: - choices = [slice(0, N - 1, None), - slice(-1, -N, -2)] - for ndim in range(maxdim + 1): - for tup in itertools.product(choices, repeat=ndim): - yield tup - - def generate_advanced_index_tuples(self, N, maxdim, many=True): - """ - Generate advanced index tuples by generating basic index tuples - and adding a single advanced index item. - """ - # (Note Numba doesn't support advanced indices with more than - # one advanced index array at the moment) - choices = list(self.generate_advanced_indices(N, many=many)) - for i in range(maxdim + 1): - for tup in self.generate_basic_index_tuples(N, maxdim - 1, many): - for adv in choices: - yield tup[:i] + (adv,) + tup[i:] - - def generate_advanced_index_tuples_with_ellipsis(self, N, maxdim, many=True): - """ - Same as generate_advanced_index_tuples(), but also insert an - ellipsis at various points. - """ - for tup in self.generate_advanced_index_tuples(N, maxdim, many): - for i in range(len(tup) + 1): - yield tup[:i] + (Ellipsis,) + tup[i:] - - def check_getitem_indices(self, arr, indices): - pyfunc = getitem_usecase - cfunc = jit(nopython=True)(pyfunc) + # Every case has exactly one array, + # Otherwise it's not fancy indexing + shape = (5, 6, 7, 8, 9, 10) + indexing_cases = [ + # Slices + Integers + (slice(4, 5), 3, np.array([0,1,3,4,2]), 1), + (3, np.array([0,1,3,4,2]), slice(None), slice(4)), + (3, np.array([[0,1,3,4,2], [0,1,2,3,2], [3,1,3,4,1]]), + slice(None), slice(4)), # multidimensional + + # Ellipsis + Integers + (Ellipsis, 1, np.array([0,1,3,4,2])), + (np.array([0,1,3,4,2]), 3, Ellipsis), + (np.array([[0,1,3,4,2], [0,1,2,3,2], [3,1,3,4,1]]), + 3, Ellipsis), # multidimensional + + # Ellipsis + Slices + Integers + (Ellipsis, 1, np.array([0,1,3,4,2]), 3, slice(1,5)), + (np.array([0,1,3,4,2]), 3, Ellipsis, slice(1,5)), + (np.array([[0,1,3,4,2], [0,1,2,3,2], [3,1,3,4,1]]), + 3, Ellipsis, slice(1,5)), # multidimensional + + # Boolean Arrays + Integers + (slice(4, 5), 3, + np.array([True, False, True, False, True, False, False]), + 1), + (3, np.array([True, False, True, False, True, False]), + slice(None), slice(4)), + ] + + rng = np.random.default_rng(1) + + def generate_random_indices(self): + N = min(self.shape) + slice_choices = [slice(None, None, None), + slice(1, N - 1, None), + slice(0, None, 2), + slice(N - 1, None, -2), + slice(-N + 1, -1, None), + slice(-1, -N, -2), + slice(0, N - 1, None), + slice(-1, -N, -2) + ] + integer_choices = list(np.arange(N)) + + indices = [] + + # Generate 20 random slice cases + for i in range(20): + array_idx = self.rng.integers(0, 5, size=15) + # Randomly select 4 slices from our list + curr_idx = self.rng.choice(slice_choices, size=4).tolist() + # Replace one of the slice with the array index + _array_idx = self.rng.choice(4) + curr_idx[_array_idx] = array_idx + indices.append(tuple(curr_idx)) + + # Generate 20 random integer cases + for i in range(20): + array_idx = self.rng.integers(0, 5, size=15) + # Randomly select 4 integers from our list + curr_idx = self.rng.choice(integer_choices, size=4).tolist() + # Replace one of the slice with the array index + _array_idx = self.rng.choice(4) + curr_idx[_array_idx] = array_idx + indices.append(tuple(curr_idx)) + + # Generate 20 random ellipsis cases + for i in range(20): + array_idx = self.rng.integers(0, 5, size=15) + # Randomly select 4 slices from our list + curr_idx = self.rng.choice(slice_choices, size=4).tolist() + # Generate two seperate random indices, replace one with + # array and second with Ellipsis + _array_idx = self.rng.choice(4, size=2, replace=False) + curr_idx[_array_idx[0]] = array_idx + curr_idx[_array_idx[1]] = Ellipsis + indices.append(tuple(curr_idx)) + + # Generate 20 random boolean cases + for i in range(20): + array_idx = self.rng.integers(0, 5, size=15) + # Randomly select 4 slices from our list + curr_idx = self.rng.choice(slice_choices, size=4).tolist() + # Replace one of the slice with the boolean array index + _array_idx = self.rng.choice(4) + bool_arr_shape = self.shape[_array_idx] + curr_idx[_array_idx] = np.array( + self.rng.choice(2, size=bool_arr_shape), + dtype=bool + ) + + indices.append(tuple(curr_idx)) + + return indices + + def check_getitem_indices(self, arr_shape, index): + def get_item(array, idx): + return array[index] + + arr = np.random.random_integers(0, 10, size=arr_shape) + get_item_numba = njit(get_item) orig = arr.copy() orig_base = arr.base or arr - for index in indices: - expected = pyfunc(arr, index) - # Sanity check: if a copy wasn't made, this wasn't advanced - # but basic indexing, and shouldn't be tested here. - assert expected.base is not orig_base - got = cfunc(arr, index) - # Note Numba may not return the same array strides and - # contiguity as Numpy - self.assertEqual(got.shape, expected.shape) - self.assertEqual(got.dtype, expected.dtype) - np.testing.assert_equal(got, expected) - # Check a copy was *really* returned by Numba - if got.size: - got.fill(42) - np.testing.assert_equal(arr, orig) - - def test_getitem_tuple(self): - # Test many variations of advanced indexing with a tuple index - N = 4 - ndim = 3 - arr = np.arange(N ** ndim).reshape((N,) * ndim).astype(np.int32) - indices = self.generate_advanced_index_tuples(N, ndim) - - self.check_getitem_indices(arr, indices) - - def test_getitem_tuple_and_ellipsis(self): - # Same, but also insert an ellipsis at a random point - N = 4 - ndim = 3 - arr = np.arange(N ** ndim).reshape((N,) * ndim).astype(np.int32) - indices = self.generate_advanced_index_tuples_with_ellipsis(N, ndim, - many=False) - - self.check_getitem_indices(arr, indices) - - def test_ellipsis_getsetitem(self): - # See https://github.com/numba/numba/issues/3225 - @jit(nopython=True) - def foo(arr, v): - arr[..., 0] = arr[..., 1] - - arr = np.arange(2) - foo(arr, 1) - self.assertEqual(arr[0], arr[1]) - - def test_getitem_array(self): - # Test advanced indexing with a single array index - N = 4 - ndim = 3 - arr = np.arange(N ** ndim).reshape((N,) * ndim).astype(np.int32) - indices = self.generate_advanced_indices(N) - self.check_getitem_indices(arr, indices) - - def check_setitem_indices(self, arr, indices): - pyfunc = setitem_usecase - cfunc = jit(nopython=True)(pyfunc) - - for index in indices: - src = arr[index] - expected = np.zeros_like(arr) - got = np.zeros_like(arr) - pyfunc(expected, index, src) - cfunc(got, index, src) - # Note Numba may not return the same array strides and - # contiguity as Numpy - self.assertEqual(got.shape, expected.shape) - self.assertEqual(got.dtype, expected.dtype) - np.testing.assert_equal(got, expected) - - def test_setitem_tuple(self): - # Test many variations of advanced indexing with a tuple index - N = 4 - ndim = 3 - arr = np.arange(N ** ndim).reshape((N,) * ndim).astype(np.int32) - indices = self.generate_advanced_index_tuples(N, ndim) - self.check_setitem_indices(arr, indices) - - def test_setitem_tuple_and_ellipsis(self): - # Same, but also insert an ellipsis at a random point - N = 4 - ndim = 3 - arr = np.arange(N ** ndim).reshape((N,) * ndim).astype(np.int32) - indices = self.generate_advanced_index_tuples_with_ellipsis(N, ndim, - many=False) - - self.check_setitem_indices(arr, indices) - - def test_setitem_array(self): - # Test advanced indexing with a single array index - N = 4 - ndim = 3 - arr = np.arange(N ** ndim).reshape((N,) * ndim).astype(np.int32) + 10 - indices = self.generate_advanced_indices(N) - self.check_setitem_indices(arr, indices) + expected = get_item(arr, index) + got = get_item_numba(arr, index) + # Sanity check: In advanced indexing, the result is always a copy. + assert expected.base is not orig_base + + # Note: Numba may not return the same array strides and + # contiguity as Numpy + self.assertEqual(got.shape, expected.shape) + self.assertEqual(got.dtype, expected.dtype) + np.testing.assert_equal(got, expected) + + # Check a copy was *really* returned by Numba + got.fill(42) + np.testing.assert_equal(arr, orig) + + def check_setitem_indices(self, arr_shape, index): + @njit + def set_item(array, idx, item): + array[idx] = item + + arr = np.random.random_integers(0, 10, size=arr_shape) + src = arr[index] + expected = np.zeros_like(arr) + got = np.zeros_like(arr) + + set_item.py_func(expected, index, src) + set_item(got, index, src) + + # Note: Numba may not return the same array strides and + # contiguity as NumPy + self.assertEqual(got.shape, expected.shape) + self.assertEqual(got.dtype, expected.dtype) + + np.testing.assert_equal(got, expected) + + def test_getitem(self): + # Cases with a combination of integers + other objects + indices = self.indexing_cases + + # Cases with permutations of either integers or objects + indices += self.generate_random_indices() + + for idx in indices: + with self.subTest(idx=idx): + self.check_getitem_indices(self.shape, idx) + + def test_setitem(self): + # Cases with a combination of integers + other objects + indices = self.indexing_cases + + # Cases with permutations of either integers or objects + indices += self.generate_random_indices() + + for idx in indices: + with self.subTest(idx=idx): + self.check_setitem_indices(self.shape, idx) + + def test_unsupported_condition_exceptions(self): + # Cases with more than one indexing array + idx = (0, 3, np.array([1, 2]), np.array([1, 2])) + with self.assertRaises(TypingError) as raises: + self.check_getitem_indices(self.shape, idx) + self.assertIn( + 'Using more than one non-scalar array index is unsupported.', + str(raises.exception) + ) + + # Cases with more than one indexing subspace + # (The subspaces here are separated by slice(None)) + idx = (0, np.array([1, 2]), slice(None), 3, 4) + with self.assertRaises(TypingError) as raises: + self.check_getitem_indices(self.shape, idx) + msg = "Using more than one indexing subspace (consecutive group " +\ + "of integer or array indices) is unsupported." + self.assertIn( + msg, + str(raises.exception) + ) def test_setitem_0d(self): + @njit + def set_item(array, idx, item): + array[idx] = item # Test setitem with a 0d-array - pyfunc = setitem_usecase - cfunc = jit(nopython=True)(pyfunc) + pyfunc = set_item.py_func + cfunc = set_item inps = [ (np.zeros(3), np.array(3.14)), @@ -199,10 +220,23 @@ def test_setitem_0d(self): cfunc(x2, 0, v) self.assertPreciseEqual(x1, x2) + def test_ellipsis_getsetitem(self): + # See https://github.com/numba/numba/issues/3225 + @njit + def foo(arr, v): + arr[..., 0] = arr[..., 1] + + arr = np.arange(2) + foo(arr, 1) + self.assertEqual(arr[0], arr[1]) + def test_np_take(self): + def np_take(array, indices): + return np.take(array, indices) + # shorter version of array.take test in test_array_methods pyfunc = np_take - cfunc = jit(nopython=True)(pyfunc) + cfunc = njit(pyfunc) def check(arr, ind): expected = pyfunc(arr, ind) @@ -243,14 +277,17 @@ def check(arr, ind): with self.assertRaises(TypingError): cfunc(A, [1.7]) + def np_take_kws(array, indices, axis): + return np.take(array, indices, axis=axis) + # check unsupported arg raises with self.assertRaises(TypingError): - take_kws = jit(nopython=True)(np_take_kws) + take_kws = njit(np_take_kws) take_kws(A, 1, 1) # check kwarg unsupported raises with self.assertRaises(TypingError): - take_kws = jit(nopython=True)(np_take_kws) + take_kws = njit(np_take_kws) take_kws(A, 1, axis=1) #exceptions leak refs diff --git a/numba/tests/test_record_dtype.py b/numba/tests/test_record_dtype.py index b76fdcaccc9..51b129ae972 100644 --- a/numba/tests/test_record_dtype.py +++ b/numba/tests/test_record_dtype.py @@ -1554,7 +1554,7 @@ def test_setitem_whole_array_error(self): nbarr2 = np.recarray(1, dtype=recordwith2darray) args = (nbarr1, nbarr2) pyfunc = record_setitem_array - errmsg = "unsupported array index type" + errmsg = "Unsupported array index type" with self.assertRaisesRegex(TypingError, errmsg): self.get_cfunc(pyfunc, tuple((typeof(arg) for arg in args)))