530 lines
17 KiB
Python
530 lines
17 KiB
Python
|
import unittest
|
||
|
from numba.tests.support import TestCase
|
||
|
|
||
|
import sys
|
||
|
import operator
|
||
|
import numpy as np
|
||
|
import numpy
|
||
|
|
||
|
from numba import jit, njit, typed
|
||
|
from numba.core import types, utils
|
||
|
from numba.core.errors import TypingError, LoweringError
|
||
|
from numba.core.types.functions import _header_lead
|
||
|
from numba.np.numpy_support import numpy_version
|
||
|
from numba.tests.support import tag, _32bit, captured_stdout
|
||
|
|
||
|
|
||
|
# deliberately imported twice for different use cases
|
||
|
|
||
|
|
||
|
PARALLEL_SUPPORTED = not _32bit
|
||
|
|
||
|
def comp_list(n):
|
||
|
l = [i for i in range(n)]
|
||
|
s = 0
|
||
|
for i in l:
|
||
|
s += i
|
||
|
return s
|
||
|
|
||
|
|
||
|
class TestListComprehension(TestCase):
|
||
|
|
||
|
def test_comp_list(self):
|
||
|
pyfunc = comp_list
|
||
|
cfunc = njit((types.intp,))(pyfunc)
|
||
|
self.assertEqual(cfunc(5), pyfunc(5))
|
||
|
self.assertEqual(cfunc(0), pyfunc(0))
|
||
|
self.assertEqual(cfunc(-1), pyfunc(-1))
|
||
|
|
||
|
def test_bulk_use_cases(self):
|
||
|
""" Tests the large number of use cases defined below """
|
||
|
|
||
|
# jitted function used in some tests
|
||
|
@jit(nopython=True)
|
||
|
def fib3(n):
|
||
|
if n < 2:
|
||
|
return n
|
||
|
return fib3(n - 1) + fib3(n - 2)
|
||
|
|
||
|
def list1(x):
|
||
|
""" Test basic list comprehension """
|
||
|
return [i for i in range(1, len(x) - 1)]
|
||
|
|
||
|
def list2(x):
|
||
|
""" Test conditional list comprehension """
|
||
|
return [y for y in x if y < 2]
|
||
|
|
||
|
def list3(x):
|
||
|
""" Test ternary list comprehension """
|
||
|
return [y if y < 2 else -1 for y in x]
|
||
|
|
||
|
def list4(x):
|
||
|
""" Test list comprehension to np.array ctor """
|
||
|
return np.array([1, 2, 3])
|
||
|
|
||
|
# expected fail, unsupported type in sequence
|
||
|
def list5(x):
|
||
|
""" Test nested list comprehension to np.array ctor """
|
||
|
return np.array([np.array([z for z in x]) for y in x])
|
||
|
|
||
|
def list6(x):
|
||
|
""" Test use of inner function in list comprehension """
|
||
|
def inner(x):
|
||
|
return x + 1
|
||
|
return [inner(z) for z in x]
|
||
|
|
||
|
def list7(x):
|
||
|
""" Test use of closure in list comprehension """
|
||
|
y = 3
|
||
|
|
||
|
def inner(x):
|
||
|
return x + y
|
||
|
return [inner(z) for z in x]
|
||
|
|
||
|
def list8(x):
|
||
|
""" Test use of list comprehension as arg to inner function """
|
||
|
l = [z + 1 for z in x]
|
||
|
|
||
|
def inner(x):
|
||
|
return x[0] + 1
|
||
|
q = inner(l)
|
||
|
return q
|
||
|
|
||
|
def list9(x):
|
||
|
""" Test use of list comprehension access in closure """
|
||
|
l = [z + 1 for z in x]
|
||
|
|
||
|
def inner(x):
|
||
|
return x[0] + l[1]
|
||
|
return inner(x)
|
||
|
|
||
|
def list10(x):
|
||
|
""" Test use of list comprehension access in closure and as arg """
|
||
|
l = [z + 1 for z in x]
|
||
|
|
||
|
def inner(x):
|
||
|
return [y + l[0] for y in x]
|
||
|
return inner(l)
|
||
|
|
||
|
def list11(x):
|
||
|
""" Test scalar array construction in list comprehension """
|
||
|
l = [np.array(z) for z in x]
|
||
|
return l
|
||
|
|
||
|
def list12(x):
|
||
|
""" Test scalar type conversion construction in list comprehension """
|
||
|
l = [np.float64(z) for z in x]
|
||
|
return l
|
||
|
|
||
|
def list13(x):
|
||
|
""" Test use of explicit numpy scalar ctor reference in list comprehension """
|
||
|
l = [numpy.float64(z) for z in x]
|
||
|
return l
|
||
|
|
||
|
def list14(x):
|
||
|
""" Test use of python scalar ctor reference in list comprehension """
|
||
|
l = [float(z) for z in x]
|
||
|
return l
|
||
|
|
||
|
def list15(x):
|
||
|
""" Test use of python scalar ctor reference in list comprehension followed by np array construction from the list"""
|
||
|
l = [float(z) for z in x]
|
||
|
return np.array(l)
|
||
|
|
||
|
def list16(x):
|
||
|
""" Test type unification from np array ctors consuming list comprehension """
|
||
|
l1 = [float(z) for z in x]
|
||
|
l2 = [z for z in x]
|
||
|
ze = np.array(l1)
|
||
|
oe = np.array(l2)
|
||
|
return ze + oe
|
||
|
|
||
|
def list17(x):
|
||
|
""" Test complex list comprehension including math calls """
|
||
|
return [(a, b, c)
|
||
|
for a in x for b in x for c in x if np.sqrt(a**2 + b**2) == c]
|
||
|
|
||
|
_OUTER_SCOPE_VAR = 9
|
||
|
|
||
|
def list18(x):
|
||
|
""" Test loop list with outer scope var as conditional"""
|
||
|
z = []
|
||
|
for i in x:
|
||
|
if i < _OUTER_SCOPE_VAR:
|
||
|
z.append(i)
|
||
|
return z
|
||
|
|
||
|
_OUTER_SCOPE_VAR = 9
|
||
|
|
||
|
def list19(x):
|
||
|
""" Test list comprehension with outer scope as conditional"""
|
||
|
return [i for i in x if i < _OUTER_SCOPE_VAR]
|
||
|
|
||
|
def list20(x):
|
||
|
""" Test return empty list """
|
||
|
return [i for i in x if i == -1000]
|
||
|
|
||
|
def list21(x):
|
||
|
""" Test call a jitted function in a list comprehension """
|
||
|
return [fib3(i) for i in x]
|
||
|
|
||
|
def list22(x):
|
||
|
""" Test create two lists comprehensions and a third walking the first two """
|
||
|
a = [y - 1 for y in x]
|
||
|
b = [y + 1 for y in x]
|
||
|
return [x for x in a for y in b if x == y]
|
||
|
|
||
|
def list23(x):
|
||
|
""" Test operation on comprehension generated list """
|
||
|
z = [y for y in x]
|
||
|
z.append(1)
|
||
|
return z
|
||
|
|
||
|
def list24(x):
|
||
|
""" Test type promotion """
|
||
|
z = [float(y) if y > 3 else y for y in x]
|
||
|
return z
|
||
|
|
||
|
def list25(x):
|
||
|
# See issue #6260. Old style inline_closure_call uses get_ir_of_code
|
||
|
# for the closure->IR transform, without SSA there's multiply
|
||
|
# defined labels, the unary negation is self referent and DCE runs
|
||
|
# eliminating the duplicated labels.
|
||
|
included = np.array([1, 2, 6, 8])
|
||
|
not_included = [i for i in range(10) if i not in list(included)]
|
||
|
return not_included
|
||
|
|
||
|
# functions to test that are expected to pass
|
||
|
f = [list1, list2, list3, list4,
|
||
|
list6, list7, list8, list9, list10, list11,
|
||
|
list12, list13, list14, list15,
|
||
|
list16, list17, list18, list19, list20,
|
||
|
list21, list22, list23, list24, list25]
|
||
|
|
||
|
var = [1, 2, 3, 4, 5]
|
||
|
for ref in f:
|
||
|
try:
|
||
|
cfunc = jit(nopython=True)(ref)
|
||
|
self.assertEqual(cfunc(var), ref(var))
|
||
|
except ValueError: # likely np array returned
|
||
|
try:
|
||
|
np.testing.assert_allclose(cfunc(var), ref(var))
|
||
|
except Exception:
|
||
|
raise
|
||
|
|
||
|
# test functions that are expected to fail
|
||
|
with self.assertRaises(TypingError) as raises:
|
||
|
cfunc = jit(nopython=True)(list5)
|
||
|
cfunc(var)
|
||
|
# TODO: we can't really assert the error message for the above
|
||
|
# Also, test_nested_array is a similar case (but without list) that works.
|
||
|
|
||
|
if sys.maxsize > 2 ** 32:
|
||
|
bits = 64
|
||
|
else:
|
||
|
bits = 32
|
||
|
|
||
|
def test_objmode_inlining(self):
|
||
|
def objmode_func(y):
|
||
|
z = object()
|
||
|
inlined = [x for x in y]
|
||
|
return inlined
|
||
|
|
||
|
cfunc = jit(forceobj=True)(objmode_func)
|
||
|
t = [1, 2, 3]
|
||
|
expected = objmode_func(t)
|
||
|
got = cfunc(t)
|
||
|
self.assertPreciseEqual(expected, got)
|
||
|
|
||
|
|
||
|
class TestArrayComprehension(unittest.TestCase):
|
||
|
|
||
|
_numba_parallel_test_ = False
|
||
|
|
||
|
def check(self, pyfunc, *args, **kwargs):
|
||
|
"""A generic check function that run both pyfunc, and jitted pyfunc,
|
||
|
and compare results."""
|
||
|
run_parallel = kwargs.get('run_parallel', False)
|
||
|
assert_allocate_list = kwargs.get('assert_allocate_list', False)
|
||
|
assert_dtype = kwargs.get('assert_dtype', False)
|
||
|
cfunc = jit(nopython=True,parallel=run_parallel)(pyfunc)
|
||
|
pyres = pyfunc(*args)
|
||
|
cres = cfunc(*args)
|
||
|
np.testing.assert_array_equal(pyres, cres)
|
||
|
if assert_dtype:
|
||
|
self.assertEqual(cres[1].dtype, assert_dtype)
|
||
|
if assert_allocate_list:
|
||
|
self.assertIn('allocate list', cfunc.inspect_llvm(cfunc.signatures[0]))
|
||
|
else:
|
||
|
self.assertNotIn('allocate list', cfunc.inspect_llvm(cfunc.signatures[0]))
|
||
|
if run_parallel:
|
||
|
self.assertIn('@do_scheduling', cfunc.inspect_llvm(cfunc.signatures[0]))
|
||
|
|
||
|
def test_comp_with_array_1(self):
|
||
|
def comp_with_array_1(n):
|
||
|
m = n * 2
|
||
|
l = np.array([i + m for i in range(n)])
|
||
|
return l
|
||
|
|
||
|
self.check(comp_with_array_1, 5)
|
||
|
if PARALLEL_SUPPORTED:
|
||
|
self.check(comp_with_array_1, 5, run_parallel=True)
|
||
|
|
||
|
def test_comp_with_array_2(self):
|
||
|
def comp_with_array_2(n, threshold):
|
||
|
A = np.arange(-n, n)
|
||
|
return np.array([ x * x if x < threshold else x * 2 for x in A ])
|
||
|
|
||
|
self.check(comp_with_array_2, 5, 0)
|
||
|
|
||
|
def test_comp_with_array_noinline(self):
|
||
|
def comp_with_array_noinline(n):
|
||
|
m = n * 2
|
||
|
l = np.array([i + m for i in range(n)])
|
||
|
return l
|
||
|
|
||
|
import numba.core.inline_closurecall as ic
|
||
|
try:
|
||
|
ic.enable_inline_arraycall = False
|
||
|
self.check(comp_with_array_noinline, 5, assert_allocate_list=True)
|
||
|
finally:
|
||
|
ic.enable_inline_arraycall = True
|
||
|
|
||
|
def test_comp_with_array_noinline_issue_6053(self):
|
||
|
def comp_with_array_noinline(n):
|
||
|
lst = [0]
|
||
|
for i in range(n):
|
||
|
lst.append(i)
|
||
|
l = np.array(lst)
|
||
|
return l
|
||
|
|
||
|
self.check(comp_with_array_noinline, 5, assert_allocate_list=True)
|
||
|
|
||
|
def test_comp_nest_with_array(self):
|
||
|
def comp_nest_with_array(n):
|
||
|
l = np.array([[i * j for j in range(n)] for i in range(n)])
|
||
|
return l
|
||
|
|
||
|
self.check(comp_nest_with_array, 5)
|
||
|
if PARALLEL_SUPPORTED:
|
||
|
self.check(comp_nest_with_array, 5, run_parallel=True)
|
||
|
|
||
|
def test_comp_nest_with_array_3(self):
|
||
|
def comp_nest_with_array_3(n):
|
||
|
l = np.array([[[i * j * k for k in range(n)] for j in range(n)] for i in range(n)])
|
||
|
return l
|
||
|
|
||
|
self.check(comp_nest_with_array_3, 5)
|
||
|
if PARALLEL_SUPPORTED:
|
||
|
self.check(comp_nest_with_array_3, 5, run_parallel=True)
|
||
|
|
||
|
def test_comp_nest_with_array_noinline(self):
|
||
|
def comp_nest_with_array_noinline(n):
|
||
|
l = np.array([[i * j for j in range(n)] for i in range(n)])
|
||
|
return l
|
||
|
|
||
|
import numba.core.inline_closurecall as ic
|
||
|
try:
|
||
|
ic.enable_inline_arraycall = False
|
||
|
self.check(comp_nest_with_array_noinline, 5,
|
||
|
assert_allocate_list=True)
|
||
|
finally:
|
||
|
ic.enable_inline_arraycall = True
|
||
|
|
||
|
def test_comp_with_array_range(self):
|
||
|
def comp_with_array_range(m, n):
|
||
|
l = np.array([i for i in range(m, n)])
|
||
|
return l
|
||
|
|
||
|
self.check(comp_with_array_range, 5, 10)
|
||
|
|
||
|
def test_comp_with_array_range_and_step(self):
|
||
|
def comp_with_array_range_and_step(m, n):
|
||
|
l = np.array([i for i in range(m, n, 2)])
|
||
|
return l
|
||
|
|
||
|
self.check(comp_with_array_range_and_step, 5, 10)
|
||
|
|
||
|
def test_comp_with_array_conditional(self):
|
||
|
def comp_with_array_conditional(n):
|
||
|
l = np.array([i for i in range(n) if i % 2 == 1])
|
||
|
return l
|
||
|
# arraycall inline would not happen when conditional is present
|
||
|
self.check(comp_with_array_conditional, 10, assert_allocate_list=True)
|
||
|
|
||
|
def test_comp_nest_with_array_conditional(self):
|
||
|
def comp_nest_with_array_conditional(n):
|
||
|
l = np.array([[i * j for j in range(n)] for i in range(n) if i % 2 == 1])
|
||
|
return l
|
||
|
self.check(comp_nest_with_array_conditional, 5,
|
||
|
assert_allocate_list=True)
|
||
|
|
||
|
@unittest.skipUnless(numpy_version < (1, 24),
|
||
|
'Setting an array element with a sequence is removed '
|
||
|
'in NumPy 1.24')
|
||
|
def test_comp_nest_with_dependency(self):
|
||
|
def comp_nest_with_dependency(n):
|
||
|
l = np.array([[i * j for j in range(i+1)] for i in range(n)])
|
||
|
return l
|
||
|
# test is expected to fail
|
||
|
with self.assertRaises(TypingError) as raises:
|
||
|
self.check(comp_nest_with_dependency, 5)
|
||
|
self.assertIn(_header_lead, str(raises.exception))
|
||
|
self.assertIn('array(undefined,', str(raises.exception))
|
||
|
|
||
|
def test_comp_unsupported_iter(self):
|
||
|
def comp_unsupported_iter():
|
||
|
val = zip([1, 2, 3], [4, 5, 6])
|
||
|
return np.array([a for a, b in val])
|
||
|
with self.assertRaises(TypingError) as raises:
|
||
|
self.check(comp_unsupported_iter)
|
||
|
self.assertIn(_header_lead, str(raises.exception))
|
||
|
self.assertIn('Unsupported iterator found in array comprehension',
|
||
|
str(raises.exception))
|
||
|
|
||
|
def test_no_array_comp(self):
|
||
|
def no_array_comp1(n):
|
||
|
l = [1,2,3,4]
|
||
|
a = np.array(l)
|
||
|
return a
|
||
|
# const 1D array is actually inlined
|
||
|
self.check(no_array_comp1, 10, assert_allocate_list=False)
|
||
|
def no_array_comp2(n):
|
||
|
l = [1,2,3,4]
|
||
|
a = np.array(l)
|
||
|
l.append(5)
|
||
|
return a
|
||
|
self.check(no_array_comp2, 10, assert_allocate_list=True)
|
||
|
|
||
|
def test_nested_array(self):
|
||
|
def nested_array(n):
|
||
|
l = np.array([ np.array([x for x in range(n)]) for y in range(n)])
|
||
|
return l
|
||
|
|
||
|
self.check(nested_array, 10)
|
||
|
|
||
|
def test_nested_array_with_const(self):
|
||
|
def nested_array(n):
|
||
|
l = np.array([ np.array([x for x in range(3)]) for y in range(4)])
|
||
|
return l
|
||
|
|
||
|
self.check(nested_array, 0)
|
||
|
|
||
|
def test_array_comp_with_iter(self):
|
||
|
def array_comp(a):
|
||
|
l = np.array([ x * x for x in a ])
|
||
|
return l
|
||
|
# with list iterator
|
||
|
l = [1,2,3,4,5]
|
||
|
self.check(array_comp, l)
|
||
|
# with array iterator
|
||
|
self.check(array_comp, np.array(l))
|
||
|
# with tuple iterator (issue #7394)
|
||
|
self.check(array_comp, tuple(l))
|
||
|
# with typed.List iterator (issue #6550)
|
||
|
self.check(array_comp, typed.List(l))
|
||
|
|
||
|
def test_array_comp_with_dtype(self):
|
||
|
def array_comp(n):
|
||
|
l = np.array([i for i in range(n)], dtype=np.complex64)
|
||
|
return l
|
||
|
|
||
|
self.check(array_comp, 10, assert_dtype=np.complex64)
|
||
|
|
||
|
def test_array_comp_inferred_dtype(self):
|
||
|
def array_comp(n):
|
||
|
l = np.array([i * 1j for i in range(n)])
|
||
|
return l
|
||
|
|
||
|
self.check(array_comp, 10)
|
||
|
|
||
|
def test_array_comp_inferred_dtype_nested(self):
|
||
|
def array_comp(n):
|
||
|
l = np.array([[i * j for j in range(n)] for i in range(n)])
|
||
|
return l
|
||
|
|
||
|
self.check(array_comp, 10)
|
||
|
|
||
|
def test_array_comp_inferred_dtype_nested_sum(self):
|
||
|
def array_comp(n):
|
||
|
l = np.array([[i * j for j in range(n)] for i in range(n)])
|
||
|
# checks that operations on the inferred array
|
||
|
return l
|
||
|
|
||
|
self.check(array_comp, 10)
|
||
|
|
||
|
def test_array_comp_inferred_dtype_outside_setitem(self):
|
||
|
def array_comp(n, v):
|
||
|
arr = np.array([i for i in range(n)])
|
||
|
# the following should not change the dtype
|
||
|
arr[0] = v
|
||
|
return arr
|
||
|
|
||
|
# float to int cast is valid
|
||
|
v = 1.2
|
||
|
self.check(array_comp, 10, v, assert_dtype=np.intp)
|
||
|
# complex to int cast is invalid
|
||
|
with self.assertRaises(TypingError) as raises:
|
||
|
cfunc = jit(nopython=True)(array_comp)
|
||
|
cfunc(10, 2.3j)
|
||
|
self.assertIn(
|
||
|
_header_lead + " Function({})".format(operator.setitem),
|
||
|
str(raises.exception),
|
||
|
)
|
||
|
self.assertIn(
|
||
|
"(array({}, 1d, C), Literal[int](0), complex128)".format(types.intp),
|
||
|
str(raises.exception),
|
||
|
)
|
||
|
|
||
|
def test_array_comp_shuffle_sideeffect(self):
|
||
|
nelem = 100
|
||
|
|
||
|
@jit(nopython=True)
|
||
|
def foo():
|
||
|
numbers = np.array([i for i in range(nelem)])
|
||
|
np.random.shuffle(numbers)
|
||
|
print(numbers)
|
||
|
|
||
|
with captured_stdout() as gotbuf:
|
||
|
foo()
|
||
|
got = gotbuf.getvalue().strip()
|
||
|
|
||
|
with captured_stdout() as expectbuf:
|
||
|
print(np.array([i for i in range(nelem)]))
|
||
|
expect = expectbuf.getvalue().strip()
|
||
|
|
||
|
# For a large enough array, the chances of shuffle to not move any
|
||
|
# element is tiny enough.
|
||
|
self.assertNotEqual(got, expect)
|
||
|
self.assertRegex(got, r'\[(\s*\d+)+\]')
|
||
|
|
||
|
def test_empty_list_not_removed(self):
|
||
|
# see issue #3724
|
||
|
def f(x):
|
||
|
t = []
|
||
|
myList = np.array([1])
|
||
|
a = np.random.choice(myList, 1)
|
||
|
t.append(x + a)
|
||
|
return a
|
||
|
self.check(f, 5, assert_allocate_list=True)
|
||
|
|
||
|
def test_reuse_of_array_var(self):
|
||
|
""" Test issue 3742 """
|
||
|
# redefinition of z breaks array comp as there's multiple defn
|
||
|
def foo(n):
|
||
|
# doesn't matter where this is in the code, it's just to ensure a
|
||
|
# `make_function` opcode exists
|
||
|
[i for i in range(1)]
|
||
|
z = np.empty(n)
|
||
|
for i in range(n):
|
||
|
z = np.zeros(n)
|
||
|
z[i] = i # write is required to trip the bug
|
||
|
|
||
|
return z
|
||
|
|
||
|
self.check(foo, 10, assert_allocate_list=True)
|
||
|
|
||
|
|
||
|
if __name__ == '__main__':
|
||
|
unittest.main()
|