import decimal import itertools import numpy as np import unittest from numba import jit, njit, typeof from numba.core import utils, types, errors from numba.tests.support import TestCase, tag from numba.core.typing import arraydecl from numba.core.types import intp, ellipsis, slice2_type, slice3_type enable_pyobj_flags = {'forceobj': True} Noflags = {'nopython': True} def slicing_1d_usecase(a, start, stop, step): return a[start:stop:step] def slicing_1d_usecase2(a, start, stop, step): b = a[start:stop:step] total = 0 for i in range(b.shape[0]): total += b[i] * (i + 1) return total def slicing_1d_usecase3(a, start, stop): b = a[start:stop] total = 0 for i in range(b.shape[0]): total += b[i] * (i + 1) return total def slicing_1d_usecase4(a): b = a[:] total = 0 for i in range(b.shape[0]): total += b[i] * (i + 1) return total def slicing_1d_usecase5(a, start): b = a[start:] total = 0 for i in range(b.shape[0]): total += b[i] * (i + 1) return total def slicing_1d_usecase6(a, stop): b = a[:stop] total = 0 for i in range(b.shape[0]): total += b[i] * (i + 1) return total def slicing_1d_usecase7(a, start): # Omitted stop with negative step (issue #1690) b = a[start::-2] total = 0 for i in range(b.shape[0]): total += b[i] * (i + 1) return total def slicing_1d_usecase8(a, start): # Omitted start with negative step b = a[::-2] total = 0 for i in range(b.shape[0]): total += b[i] * (i + 1) return total def slicing_2d_usecase(a, start1, stop1, step1, start2, stop2, step2): # The index is a homogeneous tuple of slices return a[start1:stop1:step1, start2:stop2:step2] def slicing_2d_usecase3(a, start1, stop1, step1, index): # The index is a heterogeneous tuple return a[start1:stop1:step1, index] def slicing_3d_usecase(a, index0, start1, index2): b = a[index0, start1:, index2] total = 0 for i in range(b.shape[0]): total += b[i] * (i + 1) return total def slicing_3d_usecase2(a, index0, stop1, index2): b = a[index0, :stop1, index2] total = 0 for i in range(b.shape[0]): total += b[i] * (i + 1) return total def partial_1d_usecase(a, index): b = a[index] total = 0 for i in range(b.shape[0]): total += b[i] * (i + 1) return total def integer_indexing_1d_usecase(a, i): return a[i] def integer_indexing_2d_usecase(a, i1, i2): return a[i1,i2] def integer_indexing_2d_usecase2(a, i1, i2): return a[i1][i2] def ellipsis_usecase1(a, i, j): return a[i:j, ...] def ellipsis_usecase2(a, i, j): return a[..., i:j] def ellipsis_usecase3(a, i, j): return a[i, ..., j] def none_index_usecase(a): return a[None] def empty_tuple_usecase(a): return a[()] @njit def setitem_usecase(a, index, value): a[index] = value @njit def setitem_broadcast_usecase(a, value): a[:] = value def slicing_1d_usecase_set(a, b, start, stop, step): a[start:stop:step] = b return a def slicing_1d_usecase_add(a, b, start, stop): # NOTE: uses the ROT_FOUR opcode on Python 2, only on the [start:stop] # with inplace operator form. a[start:stop] += b return a def slicing_2d_usecase_set(a, b, start, stop, step, start2, stop2, step2): a[start:stop:step,start2:stop2:step2] = b return a class TestGetItem(TestCase): """ Test basic indexed load from an array (returning a view or a scalar). Note fancy indexing is tested in test_fancy_indexing. """ def test_1d_slicing(self, flags=enable_pyobj_flags): pyfunc = slicing_1d_usecase arraytype = types.Array(types.int32, 1, 'C') argtys = (arraytype, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(10, dtype='i4') for indices in [(0, 10, 1), (2, 3, 1), (10, 0, 1), (0, 10, -1), (0, 10, 2), (9, 0, -1), (-5, -2, 1), (0, -1, 1), ]: expected = pyfunc(a, *indices) self.assertPreciseEqual(cfunc(a, *indices), expected) def test_1d_slicing_npm(self): self.test_1d_slicing(flags=Noflags) def test_1d_slicing2(self, flags=enable_pyobj_flags): pyfunc = slicing_1d_usecase2 arraytype = types.Array(types.int32, 1, 'C') argtys = (arraytype, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(10, dtype='i4') args = [(0, 10, 1), (2, 3, 1), (10, 0, 1), (0, 10, -1), (0, 10, 2)] for arg in args: self.assertEqual(pyfunc(a, *arg), cfunc(a, *arg)) # Any arraytype = types.Array(types.int32, 1, 'A') argtys = (arraytype, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(20, dtype='i4')[::2] self.assertFalse(a.flags['C_CONTIGUOUS']) self.assertFalse(a.flags['F_CONTIGUOUS']) args = [(0, 10, 1), (2, 3, 1), (10, 0, 1), (0, 10, -1), (0, 10, 2)] for arg in args: self.assertEqual(pyfunc(a, *arg), cfunc(a, *arg)) def test_1d_slicing2_npm(self): self.test_1d_slicing2(flags=Noflags) def test_1d_slicing3(self, flags=enable_pyobj_flags): pyfunc = slicing_1d_usecase3 arraytype = types.Array(types.int32, 1, 'C') argtys = (arraytype, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(10, dtype='i4') args = [(3, 10), (2, 3), (10, 0), (0, 10), (5, 10)] for arg in args: self.assertEqual(pyfunc(a, *arg), cfunc(a, *arg)) # Any arraytype = types.Array(types.int32, 1, 'A') argtys = (arraytype, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(20, dtype='i4')[::2] self.assertFalse(a.flags['C_CONTIGUOUS']) self.assertFalse(a.flags['F_CONTIGUOUS']) for arg in args: self.assertEqual(pyfunc(a, *arg), cfunc(a, *arg)) def test_1d_slicing3_npm(self): self.test_1d_slicing3(flags=Noflags) def test_1d_slicing4(self, flags=enable_pyobj_flags): pyfunc = slicing_1d_usecase4 arraytype = types.Array(types.int32, 1, 'C') argtys = (arraytype,) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(10, dtype='i4') self.assertEqual(pyfunc(a), cfunc(a)) # Any arraytype = types.Array(types.int32, 1, 'A') argtys = (arraytype,) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(20, dtype='i4')[::2] self.assertFalse(a.flags['C_CONTIGUOUS']) self.assertFalse(a.flags['F_CONTIGUOUS']) self.assertEqual(pyfunc(a), cfunc(a)) def test_1d_slicing4_npm(self): self.test_1d_slicing4(flags=Noflags) def check_1d_slicing_with_arg(self, pyfunc, flags): args = list(range(-9, 10)) arraytype = types.Array(types.int32, 1, 'C') argtys = (arraytype, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(10, dtype='i4') for arg in args: self.assertEqual(pyfunc(a, arg), cfunc(a, arg)) # Any arraytype = types.Array(types.int32, 1, 'A') argtys = (arraytype, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(20, dtype='i4')[::2] self.assertFalse(a.flags['C_CONTIGUOUS']) self.assertFalse(a.flags['F_CONTIGUOUS']) for arg in args: self.assertEqual(pyfunc(a, arg), cfunc(a, arg)) def test_1d_slicing5(self, flags=enable_pyobj_flags): pyfunc = slicing_1d_usecase5 self.check_1d_slicing_with_arg(pyfunc, flags) def test_1d_slicing5_npm(self): self.test_1d_slicing5(flags=Noflags) def test_1d_slicing6(self, flags=enable_pyobj_flags): pyfunc = slicing_1d_usecase6 self.check_1d_slicing_with_arg(pyfunc, flags) def test_1d_slicing6_npm(self): self.test_1d_slicing6(flags=Noflags) def test_1d_slicing7(self, flags=enable_pyobj_flags): pyfunc = slicing_1d_usecase7 self.check_1d_slicing_with_arg(pyfunc, flags) def test_1d_slicing7_npm(self): self.test_1d_slicing7(flags=Noflags) def test_1d_slicing8(self, flags=enable_pyobj_flags): pyfunc = slicing_1d_usecase8 self.check_1d_slicing_with_arg(pyfunc, flags) def test_1d_slicing8_npm(self): self.test_1d_slicing8(flags=Noflags) def test_2d_slicing(self, flags=enable_pyobj_flags): """ arr_2d[a:b:c] """ pyfunc = slicing_1d_usecase arraytype = types.Array(types.int32, 2, 'C') argtys = (arraytype, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(100, dtype='i4').reshape(10, 10) for args in [(0, 10, 1), (2, 3, 1), (10, 0, 1), (0, 10, -1), (0, 10, 2)]: self.assertPreciseEqual(pyfunc(a, *args), cfunc(a, *args), msg="for args %s" % (args,)) def test_2d_slicing_npm(self): self.test_2d_slicing(flags=Noflags) def test_2d_slicing2(self, flags=enable_pyobj_flags): """ arr_2d[a:b:c, d:e:f] """ # C layout pyfunc = slicing_2d_usecase arraytype = types.Array(types.int32, 2, 'C') argtys = (arraytype, types.int32, types.int32, types.int32, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(100, dtype='i4').reshape(10, 10) indices = [(0, 10, 1), (2, 3, 1), (10, 0, 1), (0, 10, -1), (0, 10, 2), (10, 0, -1), (9, 0, -2), (-5, -2, 1), (0, -1, 1), ] args = [tup1 + tup2 for (tup1, tup2) in itertools.product(indices, indices)] for arg in args: expected = pyfunc(a, *arg) self.assertPreciseEqual(cfunc(a, *arg), expected) # Any layout arraytype = types.Array(types.int32, 2, 'A') argtys = (arraytype, types.int32, types.int32, types.int32, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(400, dtype='i4').reshape(20, 20)[::2, ::2] for arg in args: expected = pyfunc(a, *arg) self.assertPreciseEqual(cfunc(a, *arg), expected) def test_2d_slicing2_npm(self): self.test_2d_slicing2(flags=Noflags) def test_2d_slicing3(self, flags=enable_pyobj_flags): """ arr_2d[a:b:c, d] """ # C layout pyfunc = slicing_2d_usecase3 arraytype = types.Array(types.int32, 2, 'C') argtys = (arraytype, types.int32, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(100, dtype='i4').reshape(10, 10) args = [ (0, 10, 1, 0), (2, 3, 1, 1), (10, 0, -1, 8), (9, 0, -2, 4), (0, 10, 2, 3), (0, -1, 3, 1), ] for arg in args: expected = pyfunc(a, *arg) self.assertPreciseEqual(cfunc(a, *arg), expected) # Any layout arraytype = types.Array(types.int32, 2, 'A') argtys = (arraytype, types.int32, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(400, dtype='i4').reshape(20, 20)[::2, ::2] for arg in args: expected = pyfunc(a, *arg) self.assertPreciseEqual(cfunc(a, *arg), expected) def test_2d_slicing3_npm(self): self.test_2d_slicing3(flags=Noflags) def test_3d_slicing(self, flags=enable_pyobj_flags): # C layout pyfunc = slicing_3d_usecase arraytype = types.Array(types.int32, 3, 'C') argtys = (arraytype, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(1000, dtype='i4').reshape(10, 10, 10) args = [ (0, 9, 1), (2, 3, 1), (9, 0, 1), (0, 9, -1), (0, 9, 2), ] for arg in args: self.assertEqual(pyfunc(a, *arg), cfunc(a, *arg)) # Any layout arraytype = types.Array(types.int32, 3, 'A') argtys = (arraytype, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(2000, dtype='i4')[::2].reshape(10, 10, 10) for arg in args: self.assertEqual(pyfunc(a, *arg), cfunc(a, *arg)) def test_3d_slicing_npm(self): self.test_3d_slicing(flags=Noflags) def test_3d_slicing2(self, flags=enable_pyobj_flags): # C layout pyfunc = slicing_3d_usecase2 arraytype = types.Array(types.int32, 3, 'C') argtys = (arraytype, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(1000, dtype='i4').reshape(10, 10, 10) args = [ (0, 9, 1), (2, 3, 1), (9, 0, 1), (0, 9, -1), (0, 9, 2), ] for arg in args: self.assertEqual(pyfunc(a, *arg), cfunc(a, *arg)) # Any layout arraytype = types.Array(types.int32, 3, 'A') argtys = (arraytype, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(2000, dtype='i4')[::2].reshape(10, 10, 10) for arg in args: self.assertEqual(pyfunc(a, *arg), cfunc(a, *arg)) def test_3d_slicing2_npm(self): self.test_3d_slicing2(flags=Noflags) def test_1d_integer_indexing(self, flags=enable_pyobj_flags): # C layout pyfunc = integer_indexing_1d_usecase arraytype = types.Array(types.int32, 1, 'C') argtys = (arraytype, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(10, dtype='i4') self.assertEqual(pyfunc(a, 0), cfunc(a, 0)) self.assertEqual(pyfunc(a, 9), cfunc(a, 9)) self.assertEqual(pyfunc(a, -1), cfunc(a, -1)) # Any layout arraytype = types.Array(types.int32, 1, 'A') argtys = (arraytype, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(10, dtype='i4')[::2] self.assertFalse(a.flags['C_CONTIGUOUS']) self.assertFalse(a.flags['F_CONTIGUOUS']) self.assertEqual(pyfunc(a, 0), cfunc(a, 0)) self.assertEqual(pyfunc(a, 2), cfunc(a, 2)) self.assertEqual(pyfunc(a, -1), cfunc(a, -1)) # Using a 0-d array as integer index arraytype = types.Array(types.int32, 1, 'C') indextype = types.Array(types.int16, 0, 'C') argtys = (arraytype, indextype) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(3, 13, dtype=np.int32) for i in (0, 9, -2): idx = np.array(i).astype(np.int16) assert idx.ndim == 0 self.assertEqual(pyfunc(a, idx), cfunc(a, idx)) def test_1d_integer_indexing_npm(self): self.test_1d_integer_indexing(flags=Noflags) def test_integer_indexing_1d_for_2d(self, flags=enable_pyobj_flags): # Test partial (1d) indexing of a 2d array pyfunc = integer_indexing_1d_usecase arraytype = types.Array(types.int32, 2, 'C') argtys = (arraytype, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(100, dtype='i4').reshape(10, 10) self.assertPreciseEqual(pyfunc(a, 0), cfunc(a, 0)) self.assertPreciseEqual(pyfunc(a, 9), cfunc(a, 9)) self.assertPreciseEqual(pyfunc(a, -1), cfunc(a, -1)) arraytype = types.Array(types.int32, 2, 'A') argtys = (arraytype, types.int32) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(20, dtype='i4').reshape(5, 4)[::2] self.assertPreciseEqual(pyfunc(a, 0), cfunc(a, 0)) def test_integer_indexing_1d_for_2d_npm(self): self.test_integer_indexing_1d_for_2d(flags=Noflags) def test_2d_integer_indexing(self, flags=enable_pyobj_flags, pyfunc=integer_indexing_2d_usecase): # C layout a = np.arange(100, dtype='i4').reshape(10, 10) arraytype = types.Array(types.int32, 2, 'C') argtys = (arraytype, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) self.assertEqual(pyfunc(a, 0, 3), cfunc(a, 0, 3)) self.assertEqual(pyfunc(a, 9, 9), cfunc(a, 9, 9)) self.assertEqual(pyfunc(a, -2, -1), cfunc(a, -2, -1)) # Any layout a = np.arange(100, dtype='i4').reshape(10, 10)[::2, ::2] self.assertFalse(a.flags['C_CONTIGUOUS']) self.assertFalse(a.flags['F_CONTIGUOUS']) arraytype = types.Array(types.int32, 2, 'A') argtys = (arraytype, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) self.assertEqual(pyfunc(a, 0, 1), cfunc(a, 0, 1)) self.assertEqual(pyfunc(a, 2, 2), cfunc(a, 2, 2)) self.assertEqual(pyfunc(a, -2, -1), cfunc(a, -2, -1)) # With 0-d arrays as integer indices a = np.arange(100, dtype='i4').reshape(10, 10) arraytype = types.Array(types.int32, 2, 'C') indextype = types.Array(types.int32, 0, 'C') argtys = (arraytype, indextype, indextype) cfunc = jit(argtys, **flags)(pyfunc) for i, j in [(0, 3), (8, 9), (-2, -1)]: i = np.array(i).astype(np.int32) j = np.array(j).astype(np.int32) self.assertEqual(pyfunc(a, i, j), cfunc(a, i, j)) def test_2d_integer_indexing_npm(self): self.test_2d_integer_indexing(flags=Noflags) def test_2d_integer_indexing2(self): self.test_2d_integer_indexing(pyfunc=integer_indexing_2d_usecase2) self.test_2d_integer_indexing(flags=Noflags, pyfunc=integer_indexing_2d_usecase2) def test_2d_integer_indexing_via_call(self): @njit def index1(X, i0): return X[i0] @njit def index2(X, i0, i1): return index1(X[i0], i1) a = np.arange(10).reshape(2, 5) self.assertEqual(index2(a, 0, 0), a[0][0]) self.assertEqual(index2(a, 1, 1), a[1][1]) self.assertEqual(index2(a, -1, -1), a[-1][-1]) def test_2d_float_indexing(self, flags=enable_pyobj_flags): a = np.arange(100, dtype='i4').reshape(10, 10) pyfunc = integer_indexing_2d_usecase arraytype = types.Array(types.int32, 2, 'C') argtys = (arraytype, types.float32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) self.assertEqual(pyfunc(a, 0, 0), cfunc(a, 0, 0)) self.assertEqual(pyfunc(a, 9, 9), cfunc(a, 9, 9)) self.assertEqual(pyfunc(a, -1, -1), cfunc(a, -1, -1)) def test_partial_1d_indexing(self, flags=enable_pyobj_flags): pyfunc = partial_1d_usecase def check(arr, arraytype): argtys = (arraytype, types.int32) cfunc = jit(argtys, **flags)(pyfunc) self.assertEqual(pyfunc(arr, 0), cfunc(arr, 0)) n = arr.shape[0] - 1 self.assertEqual(pyfunc(arr, n), cfunc(arr, n)) self.assertEqual(pyfunc(arr, -1), cfunc(arr, -1)) a = np.arange(12, dtype='i4').reshape((4, 3)) arraytype = types.Array(types.int32, 2, 'C') check(a, arraytype) a = np.arange(12, dtype='i4').reshape((3, 4)).T arraytype = types.Array(types.int32, 2, 'F') check(a, arraytype) a = np.arange(12, dtype='i4').reshape((3, 4))[::2] arraytype = types.Array(types.int32, 2, 'A') check(a, arraytype) def check_ellipsis(self, pyfunc, flags): def compile_func(arr): argtys = (typeof(arr), types.intp, types.intp) return jit(argtys, **flags)(pyfunc) def run(a): bounds = (0, 1, 2, -1, -2) cfunc = compile_func(a) for i, j in itertools.product(bounds, bounds): x = cfunc(a, i, j) np.testing.assert_equal(pyfunc(a, i, j), cfunc(a, i, j)) run(np.arange(16, dtype='i4').reshape(4, 4)) run(np.arange(27, dtype='i4').reshape(3, 3, 3)) def test_ellipsis1(self, flags=enable_pyobj_flags): self.check_ellipsis(ellipsis_usecase1, flags) def test_ellipsis1_npm(self): self.test_ellipsis1(flags=Noflags) def test_ellipsis2(self, flags=enable_pyobj_flags): self.check_ellipsis(ellipsis_usecase2, flags) def test_ellipsis2_npm(self): self.test_ellipsis2(flags=Noflags) def test_ellipsis3(self, flags=enable_pyobj_flags): self.check_ellipsis(ellipsis_usecase3, flags) def test_ellipsis3_npm(self): self.test_ellipsis3(flags=Noflags) def test_ellipsis_issue1498(self): # This is an issue due to incorrect layout inferred for when # ellpsis is used and ndenumerate is specializing on the layout. @njit def udt(arr): out = np.zeros_like(arr) i = 0 for index, val in np.ndenumerate(arr[..., i]): out[index][i] = val return out py_func = udt.py_func outersize = 4 innersize = 4 arr = np.arange(outersize * innersize).reshape(outersize, innersize) got = udt(arr) expected = py_func(arr) np.testing.assert_equal(got, expected) def test_ellipsis_issue1499(self): # This tests an issue when ndarray.__getitem__ recv a tuple of # constants. The lowering is mishandling the constant value creation. @njit def udt(arr): return arr[..., 0] arr = np.arange(3) got = udt(arr) expected = udt.py_func(arr) np.testing.assert_equal(got, expected) def test_none_index(self, flags=enable_pyobj_flags): pyfunc = none_index_usecase arraytype = types.Array(types.int32, 2, 'C') # TODO should be enable to handle this in NoPython mode argtys = (arraytype,) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(100, dtype='i4').reshape(10, 10) self.assertPreciseEqual(pyfunc(a), cfunc(a)) def test_none_index_npm(self): with self.assertTypingError(): self.test_none_index(flags=Noflags) def test_empty_tuple_indexing(self, flags=enable_pyobj_flags): pyfunc = empty_tuple_usecase arraytype = types.Array(types.int32, 0, 'C') argtys = (arraytype,) cfunc = jit(argtys, **flags)(pyfunc) a = np.arange(1, dtype='i4').reshape(()) self.assertPreciseEqual(pyfunc(a), cfunc(a)) def test_empty_tuple_indexing_npm(self): self.test_empty_tuple_indexing(flags=Noflags) class TestSetItem(TestCase): """ Test basic indexed store into an array. Note fancy indexing is tested in test_fancy_indexing. """ def test_conversion_setitem(self, flags=enable_pyobj_flags): """ this used to work, and was used in one of the tutorials """ from numba import jit def pyfunc(array): for index in range(len(array)): array[index] = index % decimal.Decimal(100) cfunc = jit("void(i8[:])", **flags)(pyfunc) udt = np.arange(100, dtype='i1') control = udt.copy() pyfunc(control) cfunc(udt) self.assertPreciseEqual(udt, control) def test_1d_slicing_set(self, flags=enable_pyobj_flags): """ 1d to 1d slice assignment """ pyfunc = slicing_1d_usecase_set # Note heterogeneous types for the source and destination arrays # (int16[:] -> int32[:]) dest_type = types.Array(types.int32, 1, 'C') src_type = types.Array(types.int16, 1, 'A') argtys = (dest_type, src_type, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) N = 10 arg = np.arange(N, dtype='i2') + 40 bounds = [0, 2, N - 2, N, N + 1, N + 3, -2, -N + 2, -N, -N - 1, -N - 3] def make_dest(): return np.zeros_like(arg, dtype='i4') for start, stop in itertools.product(bounds, bounds): for step in (1, 2, -1, -2): args = start, stop, step index = slice(*args) pyleft = pyfunc(make_dest(), arg[index], *args) cleft = cfunc(make_dest(), arg[index], *args) self.assertPreciseEqual(pyleft, cleft) # Mismatching input size and slice length with self.assertRaises(ValueError): cfunc(np.zeros_like(arg, dtype=np.int32), arg, 0, 0, 1) def check_1d_slicing_set_sequence(self, flags, seqty, seq): """ Generic sequence to 1d slice assignment """ pyfunc = slicing_1d_usecase_set dest_type = types.Array(types.int32, 1, 'C') argtys = (dest_type, seqty, types.int32, types.int32, types.int32) # This emulates the use of `compile_result`. The args that are passed # into this checking function are not as advertised in argtys and # implicit casting is required. cfunc = jit(argtys, **flags)(pyfunc).overloads[argtys].entry_point N = 10 k = len(seq) arg = np.arange(N, dtype=np.int32) args = (seq, 1, -N + k + 1, 1) expected = pyfunc(arg.copy(), *args) got = cfunc(arg.copy(), *args) self.assertPreciseEqual(expected, got) args = (seq, 1, -N + k, 1) with self.assertRaises(ValueError) as raises: cfunc(arg.copy(), *args) def test_1d_slicing_set_tuple(self, flags=enable_pyobj_flags): """ Tuple to 1d slice assignment """ self.check_1d_slicing_set_sequence( flags, types.UniTuple(types.int16, 2), (8, -42)) def test_1d_slicing_set_list(self, flags=enable_pyobj_flags): """ List to 1d slice assignment """ self.check_1d_slicing_set_sequence( flags, types.List(types.int16), [8, -42]) def test_1d_slicing_broadcast(self, flags=enable_pyobj_flags): """ scalar to 1d slice assignment """ pyfunc = slicing_1d_usecase_set arraytype = types.Array(types.int32, 1, 'C') # Note heterogeneous types for the source scalar and the destination # array (int16 -> int32[:]) argtys = (arraytype, types.int16, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) N = 10 arg = np.arange(N, dtype='i4') val = 42 bounds = [0, 2, N - 2, N, N + 1, N + 3, -2, -N + 2, -N, -N - 1, -N - 3] for start, stop in itertools.product(bounds, bounds): for step in (1, 2, -1, -2): args = val, start, stop, step pyleft = pyfunc(arg.copy(), *args) cleft = cfunc(arg.copy(), *args) self.assertPreciseEqual(pyleft, cleft) def test_1d_slicing_add(self, flags=enable_pyobj_flags): pyfunc = slicing_1d_usecase_add arraytype = types.Array(types.int32, 1, 'C') argtys = (arraytype, arraytype, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) arg = np.arange(10, dtype='i4') for test in ((0, 10), (2, 5)): pyleft = pyfunc(np.zeros_like(arg), arg[slice(*test)], *test) cleft = cfunc(np.zeros_like(arg), arg[slice(*test)], *test) self.assertPreciseEqual(pyleft, cleft) def test_1d_slicing_set_npm(self): self.test_1d_slicing_set(flags=Noflags) def test_1d_slicing_set_list_npm(self): self.test_1d_slicing_set_list(flags=Noflags) def test_1d_slicing_set_tuple_npm(self): self.test_1d_slicing_set_tuple(flags=Noflags) def test_1d_slicing_broadcast_npm(self): self.test_1d_slicing_broadcast(flags=Noflags) def test_1d_slicing_add_npm(self): self.test_1d_slicing_add(flags=Noflags) def test_2d_slicing_set(self, flags=enable_pyobj_flags): """ 2d to 2d slice assignment """ pyfunc = slicing_2d_usecase_set arraytype = types.Array(types.int32, 2, 'A') argtys = (arraytype, arraytype, types.int32, types.int32, types.int32, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) arg = np.arange(10*10, dtype='i4').reshape(10,10) tests = [ (0, 10, 1, 0, 10, 1), (2, 3, 1, 2, 3, 1), (10, 0, 1, 10, 0, 1), (0, 10, -1, 0, 10, -1), (0, 10, 2, 0, 10, 2), ] for test in tests: pyleft = pyfunc(np.zeros_like(arg), arg[slice(*test[0:3]), slice(*test[3:6])], *test) cleft = cfunc(np.zeros_like(arg), arg[slice(*test[0:3]), slice(*test[3:6])], *test) self.assertPreciseEqual(cleft, pyleft) def test_2d_slicing_broadcast(self, flags=enable_pyobj_flags): """ scalar to 2d slice assignment """ pyfunc = slicing_2d_usecase_set arraytype = types.Array(types.int32, 2, 'C') # Note heterogeneous types for the source scalar and the destination # array (int16 -> int32[:]) argtys = (arraytype, types.int16, types.int32, types.int32, types.int32, types.int32, types.int32, types.int32) cfunc = jit(argtys, **flags)(pyfunc) arg = np.arange(10*10, dtype='i4').reshape(10,10) val = 42 tests = [ (0, 10, 1, 0, 10, 1), (2, 3, 1, 2, 3, 1), (10, 0, 1, 10, 0, 1), (0, 10, -1, 0, 10, -1), (0, 10, 2, 0, 10, 2), ] for test in tests: pyleft = pyfunc(arg.copy(), val, *test) cleft = cfunc(arg.copy(), val, *test) self.assertPreciseEqual(cleft, pyleft) def test_2d_slicing_set_npm(self): self.test_2d_slicing_set(flags=Noflags) def test_2d_slicing_broadcast_npm(self): self.test_2d_slicing_broadcast(flags=Noflags) def test_setitem(self): """ scalar indexed assignment """ arr = np.arange(5) setitem_usecase(arr, 1, 42) self.assertEqual(arr.tolist(), [0, 42, 2, 3, 4]) # Using a 0-d array as scalar index setitem_usecase(arr, np.array(3).astype(np.uint16), 8) self.assertEqual(arr.tolist(), [0, 42, 2, 8, 4]) # Scalar Broadcasting arr = np.arange(9).reshape(3, 3) setitem_usecase(arr, 1, 42) self.assertEqual(arr.tolist(), [[0, 1, 2], [42, 42, 42], [6, 7, 8]]) def test_setitem_broadcast(self): """ broadcasted array assignment """ # Scalar Broadcasting dst = np.arange(5) setitem_broadcast_usecase(dst, 42) self.assertEqual(dst.tolist(), [42] * 5) # 1D -> 2D Array Broadcasting dst = np.arange(6).reshape(2, 3) setitem_broadcast_usecase(dst, np.arange(1, 4)) self.assertEqual(dst.tolist(), [[1, 2, 3], [1, 2, 3]]) # 2D -> 2D Array Broadcasting dst = np.arange(6).reshape(2, 3) setitem_broadcast_usecase(dst, np.arange(1, 4).reshape(1, 3)) self.assertEqual(dst.tolist(), [[1, 2, 3], [1, 2, 3]]) # 2D -> 4D Array Broadcasting dst = np.arange(12).reshape(2, 1, 2, 3) setitem_broadcast_usecase(dst, np.arange(1, 4).reshape(1, 3)) inner2 = [[1, 2, 3], [1, 2, 3]] self.assertEqual(dst.tolist(), [[inner2]] * 2) # 2D -> 1D Array Broadcasting dst = np.arange(5) setitem_broadcast_usecase(dst, np.arange(1, 6).reshape(1, 5)) self.assertEqual(dst.tolist(), [1, 2, 3, 4, 5]) # 4D -> 2D Array Broadcasting dst = np.arange(6).reshape(2, 3) setitem_broadcast_usecase(dst, np.arange(1, 1 + dst.size).reshape(1, 1, 2, 3)) self.assertEqual(dst.tolist(), [[1, 2, 3], [4, 5, 6]]) def test_setitem_broadcast_error(self): # higher dim assigned into lower dim # 2D -> 1D dst = np.arange(5) src = np.arange(10).reshape(2, 5) with self.assertRaises(ValueError) as raises: setitem_broadcast_usecase(dst, src) errmsg = str(raises.exception) self.assertEqual('cannot broadcast source array for assignment', errmsg) # 3D -> 2D dst = np.arange(5).reshape(1, 5) src = np.arange(10).reshape(1, 2, 5) with self.assertRaises(ValueError) as raises: setitem_broadcast_usecase(dst, src) errmsg = str(raises.exception) self.assertEqual('cannot assign slice from input of different size', errmsg) # lower to higher # 1D -> 2D dst = np.arange(10).reshape(2, 5) src = np.arange(4) with self.assertRaises(ValueError) as raises: setitem_broadcast_usecase(dst, src) errmsg = str(raises.exception) self.assertEqual('cannot assign slice from input of different size', errmsg) def test_slicing_1d_broadcast(self): # 1D -> 2D sliced (1) dst = np.arange(6).reshape(3, 2) src = np.arange(1, 3) slicing_1d_usecase_set(dst, src, 0, 2, 1) self.assertEqual(dst.tolist(), [[1, 2], [1, 2], [4, 5]]) # 1D -> 2D sliced (2) dst = np.arange(6).reshape(3, 2) src = np.arange(1, 3) slicing_1d_usecase_set(dst, src, 0, None, 2) self.assertEqual(dst.tolist(), [[1, 2], [2, 3], [1, 2]]) # 2D -> 2D sliced (3) dst = np.arange(6).reshape(3, 2) src = np.arange(1, 5).reshape(2, 2) slicing_1d_usecase_set(dst, src, None, 2, 1) self.assertEqual(dst.tolist(), [[1, 2], [3, 4], [4, 5]]) def test_setitem_readonly(self): arr = np.arange(5) arr.flags.writeable = False with self.assertRaises((TypeError, errors.TypingError)) as raises: setitem_usecase(arr, 1, 42) self.assertIn("Cannot modify readonly array of type:", str(raises.exception)) class TestTyping(TestCase): """ Check typing of basic indexing operations """ def test_layout(self): """ Check an appropriate layout is inferred for the result of array indexing. """ func = arraydecl.get_array_index_type cty = types.Array(types.float64, 3, 'C') fty = types.Array(types.float64, 3, 'F') aty = types.Array(types.float64, 3, 'A') indices = [ # Tuples of (indexing arguments, keeps "C" layout, keeps "F" layout) ((), True, True), ((ellipsis,), True, True), # Indexing from the left => can sometimes keep "C" layout ((intp,), True, False), ((slice2_type,), True, False), ((intp, slice2_type), True, False), ((slice2_type, intp), False, False), ((slice2_type, slice2_type), False, False), # Strided slices = > "A" layout ((intp, slice3_type), False, False), ((slice3_type,), False, False), # Indexing from the right => can sometimes keep "F" layout ((ellipsis, intp,), False, True), ((ellipsis, slice2_type,), False, True), ((ellipsis, intp, slice2_type,), False, False), ((ellipsis, slice2_type, intp,), False, True), ((ellipsis, slice2_type, slice2_type,), False, False), # Strided slices = > "A" layout ((ellipsis, slice3_type,), False, False), ((ellipsis, slice3_type, intp,), False, False), # Indexing from both sides => only if all dimensions are indexed ((intp, ellipsis, intp,), False, False), ((slice2_type, ellipsis, slice2_type,), False, False), ((intp, intp, slice2_type,), True, False), ((intp, ellipsis, intp, slice2_type,), True, False), ((slice2_type, intp, intp,), False, True), ((slice2_type, intp, ellipsis, intp,), False, True), ((intp, slice2_type, intp,), False, False), # Strided slices = > "A" layout ((slice3_type, intp, intp,), False, False), ((intp, intp, slice3_type,), False, False), ] for index_tuple, keep_c, _ in indices: index = types.Tuple(index_tuple) r = func(cty, index) self.assertEqual(tuple(r.index), index_tuple) self.assertEqual(r.result.layout, 'C' if keep_c else 'A', index_tuple) self.assertFalse(r.advanced) for index_tuple, _, keep_f in indices: index = types.Tuple(index_tuple) r = func(fty, index) self.assertEqual(tuple(r.index), index_tuple) self.assertEqual(r.result.layout, 'F' if keep_f else 'A', index_tuple) self.assertFalse(r.advanced) for index_tuple, _, _ in indices: index = types.Tuple(index_tuple) r = func(aty, index) self.assertEqual(tuple(r.index), index_tuple) self.assertEqual(r.result.layout, 'A') self.assertFalse(r.advanced) if __name__ == '__main__': unittest.main()