import itertools import numpy as np import sys from collections import namedtuple from io import StringIO from numba import njit, typeof, prange from numba.core import ( types, typing, ir, bytecode, postproc, cpu, registry, utils, ) from numba.tests.support import (TestCase, tag, skip_parfors_unsupported, skip_unless_scipy) from numba.parfors.array_analysis import EquivSet, ArrayAnalysis from numba.core.compiler import Compiler, Flags, PassManager from numba.core.ir_utils import remove_dead from numba.core.untyped_passes import (ExtractByteCode, TranslateByteCode, FixupArgs, IRProcessing, DeadBranchPrune, RewriteSemanticConstants, GenericRewrites, WithLifting, PreserveIR, InlineClosureLikes) from numba.core.typed_passes import (NopythonTypeInference, AnnotateTypes, NopythonRewrites, IRLegalization) from numba.core.compiler_machinery import FunctionPass, PassManager, register_pass from numba.experimental import jitclass import unittest skip_unsupported = skip_parfors_unsupported # test class for #3700 @jitclass([('L', types.int32), ('T', types.int32)]) class ExampleClass3700(object): def __init__(self, n): self.L = n self.T = n + 1 # test value for test_global_tuple GVAL = (1.2,) GVAL2 = (3, 4) class TestEquivSet(TestCase): """ Test array_analysis.EquivSet. """ def test_insert_equiv(self): s1 = EquivSet() s1.insert_equiv('a', 'b') self.assertTrue(s1.is_equiv('a', 'b')) self.assertTrue(s1.is_equiv('b', 'a')) s1.insert_equiv('c', 'd') self.assertTrue(s1.is_equiv('c', 'd')) self.assertFalse(s1.is_equiv('c', 'a')) s1.insert_equiv('a', 'c') self.assertTrue(s1.is_equiv('a', 'b', 'c', 'd')) self.assertFalse(s1.is_equiv('a', 'e')) def test_intersect(self): s1 = EquivSet() s2 = EquivSet() r = s1.intersect(s2) self.assertTrue(r.is_empty()) s1.insert_equiv('a', 'b') r = s1.intersect(s2) self.assertTrue(r.is_empty()) s2.insert_equiv('b', 'c') r = s1.intersect(s2) self.assertTrue(r.is_empty()) s2.insert_equiv('d', 'a') r = s1.intersect(s2) self.assertTrue(r.is_empty()) s1.insert_equiv('a', 'e') s2.insert_equiv('c', 'd') r = s1.intersect(s2) self.assertTrue(r.is_equiv('a', 'b')) self.assertFalse(r.is_equiv('a', 'e')) self.assertFalse(r.is_equiv('c', 'd')) @register_pass(analysis_only=False, mutates_CFG=True) class ArrayAnalysisPass(FunctionPass): _name = "array_analysis_pass" def __init__(self): FunctionPass.__init__(self) def run_pass(self, state): state.array_analysis = ArrayAnalysis(state.typingctx, state.func_ir, state.typemap, state.calltypes) state.array_analysis.run(state.func_ir.blocks) post_proc = postproc.PostProcessor(state.func_ir) post_proc.run() state.func_ir_copies.append(state.func_ir.copy()) if state.test_idempotence and len(state.func_ir_copies) > 1: state.test_idempotence(state.func_ir_copies) return False class ArrayAnalysisTester(Compiler): @classmethod def mk_pipeline(cls, args, return_type=None, flags=None, locals={}, library=None, typing_context=None, target_context=None): if not flags: flags = Flags() flags.nrt = True if typing_context is None: typing_context = registry.cpu_target.typing_context if target_context is None: target_context = registry.cpu_target.target_context return cls(typing_context, target_context, library, args, return_type, flags, locals) def compile_to_ir(self, func, test_idempotence=None): """ Populate and run compiler pipeline """ self.state.func_id = bytecode.FunctionIdentity.from_function(func) ExtractByteCode().run_pass(self.state) self.state.lifted = () self.state.lifted_from = None state = self.state state.func_ir_copies = [] state.test_idempotence = test_idempotence name = 'array_analysis_testing' pm = PassManager(name) pm.add_pass(TranslateByteCode, "analyzing bytecode") pm.add_pass(FixupArgs, "fix up args") pm.add_pass(IRProcessing, "processing IR") # pre typing if not state.flags.no_rewrites: pm.add_pass(GenericRewrites, "nopython rewrites") pm.add_pass(RewriteSemanticConstants, "rewrite semantic constants") pm.add_pass(DeadBranchPrune, "dead branch pruning") pm.add_pass(InlineClosureLikes, "inline calls to locally defined closures") # typing pm.add_pass(NopythonTypeInference, "nopython frontend") if not state.flags.no_rewrites: pm.add_pass(NopythonRewrites, "nopython rewrites") # Array Analysis pass pm.add_pass(ArrayAnalysisPass, "array analysis") if test_idempotence: # Do another pass of array analysis to test idempotence pm.add_pass(ArrayAnalysisPass, "idempotence array analysis") # legalise pm.add_pass(IRLegalization, "ensure IR is legal prior to lowering") pm.add_pass(AnnotateTypes, "annotate types") # partial compile pm.finalize() pm.run(state) return state.array_analysis class TestArrayAnalysis(TestCase): def compare_ir(self, ir_list): outputs = [] for func_ir in ir_list: remove_dead(func_ir.blocks, func_ir.arg_names, func_ir) output = StringIO() func_ir.dump(file=output) outputs.append(output.getvalue()) self.assertTrue(len(set(outputs)) == 1) # assert all outputs are equal def _compile_and_test(self, fn, arg_tys, asserts=[], equivs=[], idempotent=True): """ Compile the given function and get its IR. """ test_pipeline = ArrayAnalysisTester.mk_pipeline(arg_tys) test_idempotence = self.compare_ir if idempotent else lambda x:() analysis = test_pipeline.compile_to_ir(fn, test_idempotence) if equivs: for func in equivs: # only test the equiv_set of the first block func(analysis.equiv_sets[0]) if asserts is None: self.assertTrue(self._has_no_assertcall(analysis.func_ir)) else: for func in asserts: func(analysis.func_ir, analysis.typemap) def _has_assertcall(self, func_ir, typemap, args): msg = "Sizes of {} do not match".format(', '.join(args)) for label, block in func_ir.blocks.items(): for expr in block.find_exprs(op='call'): fn = func_ir.get_definition(expr.func.name) if isinstance(fn, ir.Global) and fn.name == 'assert_equiv': typ = typemap[expr.args[0].name] if typ.literal_value.startswith(msg): return True return False def _has_shapecall(self, func_ir, x): for label, block in func_ir.blocks.items(): for expr in block.find_exprs(op='getattr'): if expr.attr == 'shape': y = func_ir.get_definition(expr.value, lhs_only=True) z = func_ir.get_definition(x, lhs_only=True) y = y.name if isinstance(y, ir.Var) else y z = z.name if isinstance(z, ir.Var) else z if y == z: return True return False def _has_no_assertcall(self, func_ir): for label, block in func_ir.blocks.items(): for expr in block.find_exprs(op='call'): fn = func_ir.get_definition(expr.func.name) if isinstance(fn, ir.Global) and fn.name == 'assert_equiv': return False return True def with_assert(self, *args): return lambda func_ir, typemap: self.assertTrue( self._has_assertcall(func_ir, typemap, args)) def without_assert(self, *args): return lambda func_ir, typemap: self.assertFalse( self._has_assertcall(func_ir, typemap, args)) def with_equiv(self, *args): def check(equiv_set): n = len(args) for i in range(n - 1): if not equiv_set.is_equiv(args[i], args[n - 1]): return False return True return lambda equiv_set: self.assertTrue(check(equiv_set)) def without_equiv(self, *args): def check(equiv_set): n = len(args) for i in range(n - 1): if equiv_set.is_equiv(args[i], args[n - 1]): return False return True return lambda equiv_set: self.assertTrue(check(equiv_set)) def with_shapecall(self, x): return lambda func_ir, s: self.assertTrue(self._has_shapecall(func_ir, x)) def without_shapecall(self, x): return lambda func_ir, s: self.assertFalse(self._has_shapecall(func_ir, x)) def test_base_cases(self): def test_0(): a = np.zeros(0) b = np.zeros(1) m = 0 n = 1 c = np.zeros((m, n)) return self._compile_and_test(test_0, (), equivs=[self.with_equiv('a', (0,)), self.with_equiv('b', (1,)), self.with_equiv('c', (0, 1))]) def test_1(n): a = np.zeros(n) b = np.zeros(n) return a + b self._compile_and_test(test_1, (types.intp,), asserts=None) def test_2(m, n): a = np.zeros(n) b = np.zeros(m) return a + b self._compile_and_test(test_2, (types.intp, types.intp), asserts=[self.with_assert('a', 'b')]) def test_3(n): a = np.zeros(n) return a + n self._compile_and_test(test_3, (types.intp,), asserts=None) def test_4(n): a = np.zeros(n) b = a + 1 c = a + 2 return a + c self._compile_and_test(test_4, (types.intp,), asserts=None) def test_5(n): a = np.zeros((n, n)) m = n b = np.zeros((m, n)) return a + b self._compile_and_test(test_5, (types.intp,), asserts=None) def test_6(m, n): a = np.zeros(n) b = np.zeros(m) d = a + b e = a - b return d + e self._compile_and_test(test_6, (types.intp, types.intp), asserts=[self.with_assert('a', 'b'), self.without_assert('d', 'e')]) def test_7(m, n): a = np.zeros(n) b = np.zeros(m) if m == 10: d = a + b else: d = a - b return d + a self._compile_and_test(test_7, (types.intp, types.intp), asserts=[self.with_assert('a', 'b'), self.without_assert('d', 'a')]) def test_8(m, n): a = np.zeros(n) b = np.zeros(m) if m == 10: d = b + a else: d = a + a return b + d self._compile_and_test(test_8, (types.intp, types.intp), asserts=[self.with_assert('b', 'a'), self.with_assert('b', 'd')]) def test_9(m): A = np.ones(m) s = 0 while m < 2: m += 1 B = np.ones(m) s += np.sum(A + B) return s self._compile_and_test(test_9, (types.intp,), asserts=[self.with_assert('A', 'B')]) def test_10(m, n): p = m - 1 q = n + 1 r = q + 1 A = np.zeros(p) B = np.zeros(q) C = np.zeros(r) D = np.zeros(m) s = np.sum(A + B) t = np.sum(C + D) return s + t self._compile_and_test(test_10, (types.intp,types.intp,), asserts=[self.with_assert('A', 'B'), self.without_assert('C', 'D')]) def test_11(): a = np.ones(5) b = np.ones(5) c = a[1:] d = b[:-1] e = len(c) f = len(d) return e == f self._compile_and_test(test_11, (), equivs=[self.with_equiv('e', 'f')]) def test_12(): a = np.ones(25).reshape((5,5)) b = np.ones(25).reshape((5,5)) c = a[1:,:] d = b[:-1,:] e = c.shape[0] f = d.shape[0] g = len(d) return e == f self._compile_and_test(test_12, (), equivs=[self.with_equiv('e', 'f', 'g')]) def test_tup_arg(T): T2 = T return T2[0] int_arr_typ = types.Array(types.intp, 1, 'C') self._compile_and_test(test_tup_arg, (types.Tuple((int_arr_typ, int_arr_typ)),), asserts=None) def test_arr_in_tup(m): A = np.ones(m) S = (A,) B = np.ones(len(S[0])) return B self._compile_and_test(test_arr_in_tup, (types.intp,), equivs=[self.with_equiv("A", "B")]) T = namedtuple("T", ['a','b']) def test_namedtuple(n): r = T(n, n) return r[0] self._compile_and_test(test_namedtuple, (types.intp,), equivs=[self.with_equiv('r', ('n', 'n'))],) # np.where is tricky since it returns tuple of arrays def test_np_where_tup_return(A): c = np.where(A) return len(c[0]) self._compile_and_test(test_np_where_tup_return, (types.Array(types.intp, 1, 'C'),), asserts=None) def test_shape(A): (m, n) = A.shape B = np.ones((m, n)) return A + B self._compile_and_test(test_shape, (types.Array(types.intp, 2, 'C'),), asserts=None) def test_cond(l, m, n): A = np.ones(l) B = np.ones(m) C = np.ones(n) if l == m: r = np.sum(A + B) else: r = 0 if m != n: s = 0 else: s = np.sum(B + C) t = 0 if l == m: if m == n: t = np.sum(A + B + C) return r + s + t self._compile_and_test(test_cond, (types.intp, types.intp, types.intp), asserts=None) def test_assert_1(m, n): assert(m == n) A = np.ones(m) B = np.ones(n) return np.sum(A + B) self._compile_and_test(test_assert_1, (types.intp, types.intp), asserts=None) def test_assert_2(A, B): assert(A.shape == B.shape) return np.sum(A + B) self._compile_and_test(test_assert_2, (types.Array(types.intp, 1, 'C'), types.Array(types.intp, 1, 'C'),), asserts=None) self._compile_and_test(test_assert_2, (types.Array(types.intp, 2, 'C'), types.Array(types.intp, 2, 'C'),), asserts=None) # expected failure with self.assertRaises(AssertionError) as raises: self._compile_and_test(test_assert_2, (types.Array(types.intp, 1, 'C'), types.Array(types.intp, 2, 'C'),), asserts=None) msg = "Dimension mismatch" self.assertIn(msg, str(raises.exception)) def test_stencilcall(self): from numba.stencils.stencil import stencil @stencil def kernel_1(a): return 0.25 * (a[0,1] + a[1,0] + a[0,-1] + a[-1,0]) def test_1(n): a = np.ones((n,n)) b = kernel_1(a) return a + b self._compile_and_test(test_1, (types.intp,), equivs=[self.with_equiv('a', 'b')], asserts=[self.without_assert('a', 'b')]) def test_2(n): a = np.ones((n,n)) b = np.ones((n+1,n+1)) kernel_1(a, out=b) return a self._compile_and_test(test_2, (types.intp,), equivs=[self.without_equiv('a', 'b')]) @stencil(standard_indexing=('c',)) def kernel_2(a, b, c): return a[0,1,0] + b[0,-1,0] + c[0] def test_3(n): a = np.arange(64).reshape(4,8,2) b = np.arange(64).reshape(n,8,2) u = np.zeros(1) v = kernel_2(a, b, u) return v # standard indexed arrays are not considered in size equivalence self._compile_and_test(test_3, (types.intp,), equivs=[self.with_equiv('a', 'b', 'v'), self.without_equiv('a', 'u')], asserts=[self.with_assert('a', 'b')]) def test_slice(self): def test_1(m, n): A = np.zeros(m) B = np.zeros(n) s = np.sum(A + B) C = A[1:m-1] D = B[1:n-1] t = np.sum(C + D) return s + t self._compile_and_test(test_1, (types.intp,types.intp,), asserts=[self.with_assert('A', 'B'), self.without_assert('C', 'D')], idempotent=False) def test_2(m): A = np.zeros(m) B = A[0:m-3] C = A[1:m-2] D = A[2:m-1] E = B + C return D + E self._compile_and_test(test_2, (types.intp,), asserts=[self.without_assert('B', 'C'), self.without_assert('D', 'E')], idempotent=False) def test_3(m): A = np.zeros((m,m)) B = A[0:m-2,0:m-2] C = A[1:m-1,1:m-1] E = B + C return E self._compile_and_test(test_3, (types.intp,), asserts=[self.without_assert('B', 'C')], idempotent=False) def test_4(m): A = np.zeros((m,m)) B = A[0:m-2,:] C = A[1:m-1,:] E = B + C return E self._compile_and_test(test_4, (types.intp,), asserts=[self.without_assert('B', 'C')], idempotent=False) def test_5(m,n): A = np.zeros(m) B = np.zeros(m) B[0:m-2] = A[1:m-1] C = np.zeros(n) D = A[1:m-1] C[0:n-2] = D # B and C are not necessarily of the same size because we can't # derive m == n from (m-2) % m == (n-2) % n return B + C self._compile_and_test(test_5, (types.intp,types.intp), asserts=[self.without_assert('B', 'A'), self.with_assert('C', 'D'), self.with_assert('B', 'C')], idempotent=False) def test_6(m): A = np.zeros((m,m)) B = A[0:m-2,:-1] C = A[1:m-1,:-1] E = B + C return E self._compile_and_test(test_6, (types.intp,), asserts=[self.without_assert('B', 'C')], idempotent=False) def test_7(m): A = np.zeros((m,m)) B = A[0:m-2,-3:-1] C = A[1:m-1,-4:-2] E = B + C return E self._compile_and_test(test_7, (types.intp,), asserts=[self.with_assert('B', 'C')], idempotent=False) def test_8(m): A = np.zeros((m,m)) B = A[:m-2,0:] C = A[1:-1,:] E = B + C return E self._compile_and_test(test_8, (types.intp,), asserts=[self.without_assert('B', 'C')], idempotent=False) def test_9(m): # issues #3461 and #3554, checks equivalence on empty slices # and across binop A = np.zeros((m)) B = A[:0] # B = array([], dtype=int64) C = A[1:] D = A[:-1:-1] # D = array([], dtype=int64) E = B + D F = E F += 1 # F = array([], dtype=int64) return A, C, F self._compile_and_test(test_9, (types.intp,), equivs=[self.without_equiv('B', 'C'), self.with_equiv('A', 'm'), self.with_equiv('B', 'D'), self.with_equiv('F', 'D'),], idempotent=False) @skip_unless_scipy def test_numpy_calls(self): def test_zeros(n): a = np.zeros(n) b = np.zeros((n, n)) c = np.zeros(shape=(n, n)) self._compile_and_test(test_zeros, (types.intp,), equivs=[self.with_equiv('a', 'n'), self.with_equiv('b', ('n', 'n')), self.with_equiv('b', 'c')]) def test_0d_array(n): a = np.array(1) b = np.ones(2) return a + b self._compile_and_test(test_0d_array, (types.intp,), equivs=[self.without_equiv('a', 'b')], asserts=[self.without_shapecall('a')]) def test_ones(n): a = np.ones(n) b = np.ones((n, n)) c = np.ones(shape=(n, n)) self._compile_and_test(test_ones, (types.intp,), equivs=[self.with_equiv('a', 'n'), self.with_equiv('b', ('n', 'n')), self.with_equiv('b', 'c')]) def test_empty(n): a = np.empty(n) b = np.empty((n, n)) c = np.empty(shape=(n, n)) self._compile_and_test(test_empty, (types.intp,), equivs=[self.with_equiv('a', 'n'), self.with_equiv('b', ('n', 'n')), self.with_equiv('b', 'c')]) def test_eye(n): a = np.eye(n) b = np.eye(N=n) c = np.eye(N=n, M=n) d = np.eye(N=n, M=n + 1) self._compile_and_test(test_eye, (types.intp,), equivs=[self.with_equiv('a', ('n', 'n')), self.with_equiv('b', ('n', 'n')), self.with_equiv('b', 'c'), self.without_equiv('b', 'd')]) def test_identity(n): a = np.identity(n) self._compile_and_test(test_identity, (types.intp,), equivs=[self.with_equiv('a', ('n', 'n'))]) def test_diag(n): a = np.identity(n) b = np.diag(a) c = np.diag(b) d = np.diag(a, k=1) self._compile_and_test(test_diag, (types.intp,), equivs=[self.with_equiv('b', ('n',)), self.with_equiv('c', ('n', 'n'))], asserts=[self.with_shapecall('d'), self.without_shapecall('c')]) def test_array_like(a): b = np.empty_like(a) c = np.zeros_like(a) d = np.ones_like(a) e = np.full_like(a, 1) f = np.asfortranarray(a) self._compile_and_test(test_array_like, (types.Array(types.intp, 2, 'C'),), equivs=[ self.with_equiv('a', 'b', 'd', 'e', 'f')], asserts=[self.with_shapecall('a'), self.without_shapecall('b')]) def test_reshape(n): a = np.ones(n * n) b = a.reshape((n, n)) return a.sum() + b.sum() self._compile_and_test(test_reshape, (types.intp,), equivs=[self.with_equiv('b', ('n', 'n'))], asserts=[self.without_shapecall('b')]) def test_transpose(m, n): a = np.ones((m, n)) b = a.T c = a.transpose() # Numba njit cannot compile explicit transpose call! # c = np.transpose(b) self._compile_and_test(test_transpose, (types.intp, types.intp), equivs=[self.with_equiv('a', ('m', 'n')), self.with_equiv('b', ('n', 'm')), self.with_equiv('c', ('n', 'm'))]) def test_transpose_3d(m, n, k): a = np.ones((m, n, k)) b = a.T c = a.transpose() d = a.transpose(2,0,1) dt = a.transpose((2,0,1)) e = a.transpose(0,2,1) et = a.transpose((0,2,1)) # Numba njit cannot compile explicit transpose call! # c = np.transpose(b) self._compile_and_test(test_transpose_3d, (types.intp, types.intp, types.intp), equivs=[self.with_equiv('a', ('m', 'n', 'k')), self.with_equiv('b', ('k', 'n', 'm')), self.with_equiv('c', ('k', 'n', 'm')), self.with_equiv('d', ('k', 'm', 'n')), self.with_equiv('dt', ('k', 'm', 'n')), self.with_equiv('e', ('m', 'k', 'n')), self.with_equiv('et', ('m', 'k', 'n'))]) def test_real_imag_attr(m, n): a = np.ones((m, n)) b = a.real c = a.imag self._compile_and_test(test_real_imag_attr, (types.intp, types.intp), equivs=[self.with_equiv('a', ('m', 'n')), self.with_equiv('b', ('m', 'n')), self.with_equiv('c', ('m', 'n')), ]) def test_random(n): a0 = np.random.rand(n) a1 = np.random.rand(n, n) b0 = np.random.randn(n) b1 = np.random.randn(n, n) c0 = np.random.ranf(n) c1 = np.random.ranf((n, n)) c2 = np.random.ranf(size=(n, n)) d0 = np.random.random_sample(n) d1 = np.random.random_sample((n, n)) d2 = np.random.random_sample(size=(n, n)) e0 = np.random.sample(n) e1 = np.random.sample((n, n)) e2 = np.random.sample(size=(n, n)) f0 = np.random.random(n) f1 = np.random.random((n, n)) f2 = np.random.random(size=(n, n)) g0 = np.random.standard_normal(n) g1 = np.random.standard_normal((n, n)) g2 = np.random.standard_normal(size=(n, n)) h0 = np.random.chisquare(10, n) h1 = np.random.chisquare(10, (n, n)) h2 = np.random.chisquare(10, size=(n, n)) i0 = np.random.weibull(10, n) i1 = np.random.weibull(10, (n, n)) i2 = np.random.weibull(10, size=(n, n)) j0 = np.random.power(10, n) j1 = np.random.power(10, (n, n)) j2 = np.random.power(10, size=(n, n)) k0 = np.random.geometric(0.1, n) k1 = np.random.geometric(0.1, (n, n)) k2 = np.random.geometric(0.1, size=(n, n)) l0 = np.random.exponential(10, n) l1 = np.random.exponential(10, (n, n)) l2 = np.random.exponential(10, size=(n, n)) m0 = np.random.poisson(10, n) m1 = np.random.poisson(10, (n, n)) m2 = np.random.poisson(10, size=(n, n)) n0 = np.random.rayleigh(10, n) n1 = np.random.rayleigh(10, (n, n)) n2 = np.random.rayleigh(10, size=(n, n)) o0 = np.random.normal(0, 1, n) o1 = np.random.normal(0, 1, (n, n)) o2 = np.random.normal(0, 1, size=(n, n)) p0 = np.random.uniform(0, 1, n) p1 = np.random.uniform(0, 1, (n, n)) p2 = np.random.uniform(0, 1, size=(n, n)) q0 = np.random.beta(0.1, 1, n) q1 = np.random.beta(0.1, 1, (n, n)) q2 = np.random.beta(0.1, 1, size=(n, n)) r0 = np.random.binomial(0, 1, n) r1 = np.random.binomial(0, 1, (n, n)) r2 = np.random.binomial(0, 1, size=(n, n)) s0 = np.random.f(0.1, 1, n) s1 = np.random.f(0.1, 1, (n, n)) s2 = np.random.f(0.1, 1, size=(n, n)) t0 = np.random.gamma(0.1, 1, n) t1 = np.random.gamma(0.1, 1, (n, n)) t2 = np.random.gamma(0.1, 1, size=(n, n)) u0 = np.random.lognormal(0, 1, n) u1 = np.random.lognormal(0, 1, (n, n)) u2 = np.random.lognormal(0, 1, size=(n, n)) v0 = np.random.laplace(0, 1, n) v1 = np.random.laplace(0, 1, (n, n)) v2 = np.random.laplace(0, 1, size=(n, n)) w0 = np.random.randint(0, 10, n) w1 = np.random.randint(0, 10, (n, n)) w2 = np.random.randint(0, 10, size=(n, n)) x0 = np.random.triangular(-3, 0, 10, n) x1 = np.random.triangular(-3, 0, 10, (n, n)) x2 = np.random.triangular(-3, 0, 10, size=(n, n)) last = ord('x') + 1 vars1d = [('n',)] + [chr(x) + '0' for x in range(ord('a'), last)] vars2d = [('n', 'n')] + [chr(x) + '1' for x in range(ord('a'), last)] vars2d += [chr(x) + '1' for x in range(ord('c'), last)] self._compile_and_test(test_random, (types.intp,), equivs=[self.with_equiv(*vars1d), self.with_equiv(*vars2d)]) def test_concatenate(m, n): a = np.ones(m) b = np.ones(n) c = np.concatenate((a, b)) d = np.ones((2, n)) e = np.ones((3, n)) f = np.concatenate((d, e)) # Numba njit cannot compile concatenate with single array! # g = np.ones((3,4,5)) # h = np.concatenate(g) i = np.ones((m, 2)) j = np.ones((m, 3)) k = np.concatenate((i, j), axis=1) l = np.ones((m, n)) o = np.ones((m, n)) p = np.concatenate((l, o)) # Numba njit cannot support list argument! # q = np.concatenate([d, e]) self._compile_and_test(test_concatenate, (types.intp, types.intp), equivs=[self.with_equiv('f', (5, 'n')), #self.with_equiv('h', (3 + 4 + 5, )), self.with_equiv('k', ('m', 5))], asserts=[self.with_shapecall('c'), self.without_shapecall('f'), self.without_shapecall('k'), self.with_shapecall('p')]) def test_vsd_stack(): k = np.ones((2,)) l = np.ones((2, 3)) o = np.ones((2, 3, 4)) p = np.vstack((k, k)) q = np.vstack((l, l)) r = np.hstack((k, k)) s = np.hstack((l, l)) t = np.dstack((k, k)) u = np.dstack((l, l)) v = np.dstack((o, o)) self._compile_and_test(test_vsd_stack, (), equivs=[self.with_equiv('p', (2, 2)), self.with_equiv('q', (4, 3)), self.with_equiv('r', (4,)), self.with_equiv('s', (2, 6)), self.with_equiv('t', (1, 2, 2)), self.with_equiv('u', (2, 3, 2)), self.with_equiv('v', (2, 3, 8)), ]) def test_stack(m, n): a = np.ones(m) b = np.ones(n) c = np.stack((a, b)) d = np.ones((m, n)) e = np.ones((m, n)) f = np.stack((d, e)) g = np.stack((d, e), axis=0) h = np.stack((d, e), axis=1) i = np.stack((d, e), axis=2) j = np.stack((d, e), axis=-1) self._compile_and_test(test_stack, (types.intp, types.intp), equivs=[self.with_equiv('m', 'n'), self.with_equiv('c', (2, 'm')), self.with_equiv( 'f', 'g', (2, 'm', 'n')), self.with_equiv( 'h', ('m', 2, 'n')), self.with_equiv( 'i', 'j', ('m', 'n', 2)), ]) def test_linspace(m, n): a = np.linspace(m, n) b = np.linspace(m, n, 10) # Numba njit does not support num keyword to linspace call! # c = np.linspace(m,n,num=10) self._compile_and_test(test_linspace, (types.float64, types.float64), equivs=[self.with_equiv('a', (50,)), self.with_equiv('b', (10,))]) def test_dot(l, m, n): a = np.dot(np.ones(1), np.ones(1)) b = np.dot(np.ones(2), np.ones((2, 3))) # Numba njit does not support higher dimensional inputs #c = np.dot(np.ones(2),np.ones((3,2,4))) #d = np.dot(np.ones(2),np.ones((3,5,2,4))) e = np.dot(np.ones((1, 2)), np.ones(2,)) #f = np.dot(np.ones((1,2,3)),np.ones(3,)) #g = np.dot(np.ones((1,2,3,4)),np.ones(4,)) h = np.dot(np.ones((2, 3)), np.ones((3, 4))) i = np.dot(np.ones((m, n)), np.ones((n, m))) j = np.dot(np.ones((m, m)), np.ones((l, l))) self._compile_and_test(test_dot, (types.intp, types.intp, types.intp), equivs=[self.without_equiv('a', (1,)), # not array self.with_equiv('b', (3,)), self.with_equiv('e', (1,)), self.with_equiv('h', (2, 4)), self.with_equiv('i', ('m', 'm')), self.with_equiv('j', ('m', 'm')), ], asserts=[self.with_assert('m', 'l')]) def test_broadcast(m, n): a = np.ones((m, n)) b = np.ones(n) c = a + b d = np.ones((1, n)) e = a + c - d self._compile_and_test(test_broadcast, (types.intp, types.intp), equivs=[self.with_equiv('a', 'c', 'e')], asserts=None) # make sure shape of a global tuple of ints is handled properly def test_global_tuple(): a = np.ones(GVAL2) b = np.ones(GVAL2) self._compile_and_test(test_global_tuple, (), equivs=[self.with_equiv('a', 'b')], asserts=None) class TestArrayAnalysisParallelRequired(TestCase): """This is to just split out tests that need the parallel backend and therefore serialised execution. """ _numba_parallel_test_ = False @skip_unsupported def test_misc(self): @njit def swap(x, y): return(y, x) def test_bug2537(m): a = np.ones(m) b = np.ones(m) for i in range(m): a[i], b[i] = swap(a[i], b[i]) try: njit(test_bug2537, parallel=True)(10) except IndexError: self.fail("test_bug2537 raised IndexError!") @skip_unsupported def test_global_namedtuple(self): Row = namedtuple('Row', ['A']) row = Row(3) def test_impl(): rr = row res = rr.A if res == 2: res = 3 return res self.assertEqual(njit(test_impl, parallel=True)(), test_impl()) @skip_unsupported def test_array_T_issue_3700(self): def test_impl(t_obj, X): for i in prange(t_obj.T): X[i] = i return X.sum() n = 5 t_obj = ExampleClass3700(n) X1 = np.zeros(t_obj.T) X2 = np.zeros(t_obj.T) self.assertEqual( njit(test_impl, parallel=True)(t_obj, X1), test_impl(t_obj, X2)) @skip_unsupported def test_slice_shape_issue_3380(self): # these tests shouldn't throw error in array analysis def test_impl1(): a = slice(None, None) return True self.assertEqual(njit(test_impl1, parallel=True)(), test_impl1()) def test_impl2(A, a): b = a return A[b] A = np.arange(10) a = slice(None) np.testing.assert_array_equal( njit(test_impl2, parallel=True)(A, a), test_impl2(A, a)) @skip_unsupported def test_slice_dtype_issue_5056(self): # see issue 5056 @njit(parallel=True) def test_impl(data): N = data.shape[0] sums = np.zeros(N) for i in prange(N): sums[i] = np.sum(data[np.int32(0):np.int32(1)]) return sums data = np.arange(10.) np.testing.assert_array_equal(test_impl(data), test_impl.py_func(data)) @skip_unsupported def test_global_tuple(self): """make sure a global tuple with non-integer values does not cause errors (test for #6726). """ def test_impl(): d = GVAL[0] return d self.assertEqual(njit(test_impl, parallel=True)(), test_impl()) class TestArrayAnalysisInterface(TestCase): def test_analyze_op_call_interface(self): # gather _analyze_op_call_* aoc = {} for fname in dir(ArrayAnalysis): if fname.startswith('_analyze_op_call_'): aoc[fname] = getattr(ArrayAnalysis, fname) # check interface def iface_stub(self, scope, equiv_set, loc, args, kws): pass expected = utils.pysignature(iface_stub) for k, v in aoc.items(): got = utils.pysignature(v) with self.subTest(fname=k, sig=got): self.assertEqual(got, expected) @skip_unsupported def test_array_analysis_extensions(self): # Test that the `array_analysis` object in `array_analysis_extensions` # can perform analysis on the scope using `equiv_sets`. from numba.parfors.parfor import Parfor from numba.parfors import array_analysis orig_parfor = array_analysis.array_analysis_extensions[Parfor] shared = {'counter': 0} def testcode(array_analysis): # Find call node corresponding to the ``A = empty(n)`` func_ir = array_analysis.func_ir for call in func_ir.blocks[0].find_exprs('call'): callee = func_ir.get_definition(call.func) if getattr(callee, "value", None) is empty: if getattr(call.args[0], 'name', None) == 'n': break else: return variable_A = func_ir.get_assignee(call) # n must be equiv to es = array_analysis.equiv_sets[0] self.assertTrue(es.is_equiv('n', variable_A.name)) shared['counter'] += 1 def new_parfor(parfor, equiv_set, typemap, array_analysis): """Recursive array analysis for parfor nodes. """ testcode(array_analysis) # Call original return orig_parfor( parfor, equiv_set, typemap, array_analysis, ) try: # Replace the array-analysis extension for Parfor node array_analysis.array_analysis_extensions[Parfor] = new_parfor empty = np.empty # avoid scanning a getattr in the IR def f(n): A = empty(n) for i in prange(n): S = np.arange(i) A[i] = S.sum() return A + 1 got = njit(parallel=True)(f)(10) executed_count = shared['counter'] self.assertGreater(executed_count, 0) finally: # Re-install the original handler array_analysis.array_analysis_extensions[Parfor] = orig_parfor # Check normal execution expected = njit(parallel=True)(f)(10) self.assertPreciseEqual(got, expected) # Make sure we have uninstalled the handler self.assertEqual(executed_count, shared['counter']) if __name__ == '__main__': unittest.main()