1850 lines
54 KiB
Python
1850 lines
54 KiB
Python
|
import contextlib
|
||
|
|
import sys
|
||
|
|
import numpy as np
|
||
|
|
import random
|
||
|
|
import re
|
||
|
|
import threading
|
||
|
|
import gc
|
||
|
|
|
||
|
|
from numba.core.errors import TypingError
|
||
|
|
from numba import njit
|
||
|
|
from numba.core import types, utils, config
|
||
|
|
from numba.tests.support import MemoryLeakMixin, TestCase, tag, skip_if_32bit
|
||
|
|
import unittest
|
||
|
|
|
||
|
|
|
||
|
|
nrtjit = njit(_nrt=True, nogil=True)
|
||
|
|
|
||
|
|
|
||
|
|
def np_concatenate1(a, b, c):
|
||
|
|
return np.concatenate((a, b, c))
|
||
|
|
|
||
|
|
def np_concatenate2(a, b, c, axis):
|
||
|
|
return np.concatenate((a, b, c), axis=axis)
|
||
|
|
|
||
|
|
def np_stack1(a, b, c):
|
||
|
|
return np.stack((a, b, c))
|
||
|
|
|
||
|
|
def np_stack2(a, b, c, axis):
|
||
|
|
return np.stack((a, b, c), axis=axis)
|
||
|
|
|
||
|
|
def np_hstack(a, b, c):
|
||
|
|
return np.hstack((a, b, c))
|
||
|
|
|
||
|
|
def np_vstack(a, b, c):
|
||
|
|
return np.vstack((a, b, c))
|
||
|
|
|
||
|
|
def np_row_stack(a, b, c):
|
||
|
|
return np.row_stack((a, b, c))
|
||
|
|
|
||
|
|
def np_dstack(a, b, c):
|
||
|
|
return np.dstack((a, b, c))
|
||
|
|
|
||
|
|
def np_column_stack(a, b, c):
|
||
|
|
return np.column_stack((a, b, c))
|
||
|
|
|
||
|
|
|
||
|
|
class BaseTest(TestCase):
|
||
|
|
|
||
|
|
def check_outputs(self, pyfunc, argslist, exact=True):
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
for args in argslist:
|
||
|
|
expected = pyfunc(*args)
|
||
|
|
ret = cfunc(*args)
|
||
|
|
self.assertEqual(ret.size, expected.size)
|
||
|
|
self.assertEqual(ret.dtype, expected.dtype)
|
||
|
|
self.assertStridesEqual(ret, expected)
|
||
|
|
if exact:
|
||
|
|
np.testing.assert_equal(expected, ret)
|
||
|
|
else:
|
||
|
|
np.testing.assert_allclose(expected, ret)
|
||
|
|
|
||
|
|
|
||
|
|
class NrtRefCtTest(MemoryLeakMixin):
|
||
|
|
def assert_array_nrt_refct(self, arr, expect):
|
||
|
|
self.assertEqual(arr.base.refcount, expect)
|
||
|
|
|
||
|
|
|
||
|
|
class TestDynArray(NrtRefCtTest, TestCase):
|
||
|
|
|
||
|
|
def test_empty_0d(self):
|
||
|
|
@nrtjit
|
||
|
|
def foo():
|
||
|
|
arr = np.empty(())
|
||
|
|
arr[()] = 42
|
||
|
|
return arr
|
||
|
|
|
||
|
|
arr = foo()
|
||
|
|
self.assert_array_nrt_refct(arr, 1)
|
||
|
|
np.testing.assert_equal(42, arr)
|
||
|
|
self.assertEqual(arr.size, 1)
|
||
|
|
self.assertEqual(arr.shape, ())
|
||
|
|
self.assertEqual(arr.dtype, np.dtype(np.float64))
|
||
|
|
self.assertEqual(arr.strides, ())
|
||
|
|
arr.fill(123) # test writability
|
||
|
|
np.testing.assert_equal(123, arr)
|
||
|
|
del arr
|
||
|
|
|
||
|
|
def test_empty_1d(self):
|
||
|
|
@nrtjit
|
||
|
|
def foo(n):
|
||
|
|
arr = np.empty(n)
|
||
|
|
for i in range(n):
|
||
|
|
arr[i] = i
|
||
|
|
|
||
|
|
return arr
|
||
|
|
|
||
|
|
n = 3
|
||
|
|
arr = foo(n)
|
||
|
|
self.assert_array_nrt_refct(arr, 1)
|
||
|
|
np.testing.assert_equal(np.arange(n), arr)
|
||
|
|
self.assertEqual(arr.size, n)
|
||
|
|
self.assertEqual(arr.shape, (n,))
|
||
|
|
self.assertEqual(arr.dtype, np.dtype(np.float64))
|
||
|
|
self.assertEqual(arr.strides, (np.dtype(np.float64).itemsize,))
|
||
|
|
arr.fill(123) # test writability
|
||
|
|
np.testing.assert_equal(123, arr)
|
||
|
|
del arr
|
||
|
|
|
||
|
|
def test_empty_2d(self):
|
||
|
|
def pyfunc(m, n):
|
||
|
|
arr = np.empty((m, n), np.int32)
|
||
|
|
for i in range(m):
|
||
|
|
for j in range(n):
|
||
|
|
arr[i, j] = i + j
|
||
|
|
|
||
|
|
return arr
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
m = 4
|
||
|
|
n = 3
|
||
|
|
expected_arr = pyfunc(m, n)
|
||
|
|
got_arr = cfunc(m, n)
|
||
|
|
self.assert_array_nrt_refct(got_arr, 1)
|
||
|
|
np.testing.assert_equal(expected_arr, got_arr)
|
||
|
|
|
||
|
|
self.assertEqual(expected_arr.size, got_arr.size)
|
||
|
|
self.assertEqual(expected_arr.shape, got_arr.shape)
|
||
|
|
self.assertEqual(expected_arr.strides, got_arr.strides)
|
||
|
|
|
||
|
|
del got_arr
|
||
|
|
|
||
|
|
def test_empty_3d(self):
|
||
|
|
def pyfunc(m, n, p):
|
||
|
|
arr = np.empty((m, n, p), np.int32)
|
||
|
|
for i in range(m):
|
||
|
|
for j in range(n):
|
||
|
|
for k in range(p):
|
||
|
|
arr[i, j, k] = i + j + k
|
||
|
|
|
||
|
|
return arr
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
m = 4
|
||
|
|
n = 3
|
||
|
|
p = 2
|
||
|
|
expected_arr = pyfunc(m, n, p)
|
||
|
|
got_arr = cfunc(m, n, p)
|
||
|
|
self.assert_array_nrt_refct(got_arr, 1)
|
||
|
|
np.testing.assert_equal(expected_arr, got_arr)
|
||
|
|
|
||
|
|
self.assertEqual(expected_arr.size, got_arr.size)
|
||
|
|
self.assertEqual(expected_arr.shape, got_arr.shape)
|
||
|
|
self.assertEqual(expected_arr.strides, got_arr.strides)
|
||
|
|
|
||
|
|
del got_arr
|
||
|
|
|
||
|
|
def test_empty_2d_sliced(self):
|
||
|
|
def pyfunc(m, n, p):
|
||
|
|
arr = np.empty((m, n), np.int32)
|
||
|
|
for i in range(m):
|
||
|
|
for j in range(n):
|
||
|
|
arr[i, j] = i + j
|
||
|
|
|
||
|
|
return arr[p]
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
m = 4
|
||
|
|
n = 3
|
||
|
|
p = 2
|
||
|
|
expected_arr = pyfunc(m, n, p)
|
||
|
|
got_arr = cfunc(m, n, p)
|
||
|
|
self.assert_array_nrt_refct(got_arr, 1)
|
||
|
|
np.testing.assert_equal(expected_arr, got_arr)
|
||
|
|
|
||
|
|
self.assertEqual(expected_arr.size, got_arr.size)
|
||
|
|
self.assertEqual(expected_arr.shape, got_arr.shape)
|
||
|
|
self.assertEqual(expected_arr.strides, got_arr.strides)
|
||
|
|
|
||
|
|
del got_arr
|
||
|
|
|
||
|
|
def test_return_global_array(self):
|
||
|
|
y = np.ones(4, dtype=np.float32)
|
||
|
|
initrefct = sys.getrefcount(y)
|
||
|
|
|
||
|
|
def return_external_array():
|
||
|
|
return y
|
||
|
|
|
||
|
|
cfunc = nrtjit(return_external_array)
|
||
|
|
out = cfunc()
|
||
|
|
|
||
|
|
# out reference by cfunc
|
||
|
|
self.assertEqual(initrefct + 1, sys.getrefcount(y))
|
||
|
|
|
||
|
|
np.testing.assert_equal(y, out)
|
||
|
|
np.testing.assert_equal(y, np.ones(4, dtype=np.float32))
|
||
|
|
np.testing.assert_equal(out, np.ones(4, dtype=np.float32))
|
||
|
|
|
||
|
|
del out
|
||
|
|
gc.collect()
|
||
|
|
# out is only referenced by cfunc
|
||
|
|
self.assertEqual(initrefct + 1, sys.getrefcount(y))
|
||
|
|
|
||
|
|
del cfunc
|
||
|
|
gc.collect()
|
||
|
|
# y is no longer referenced by cfunc
|
||
|
|
self.assertEqual(initrefct, sys.getrefcount(y))
|
||
|
|
|
||
|
|
def test_return_global_array_sliced(self):
|
||
|
|
y = np.ones(4, dtype=np.float32)
|
||
|
|
|
||
|
|
def return_external_array():
|
||
|
|
return y[2:]
|
||
|
|
|
||
|
|
cfunc = nrtjit(return_external_array)
|
||
|
|
out = cfunc()
|
||
|
|
self.assertIsNone(out.base)
|
||
|
|
|
||
|
|
yy = y[2:]
|
||
|
|
np.testing.assert_equal(yy, out)
|
||
|
|
np.testing.assert_equal(yy, np.ones(2, dtype=np.float32))
|
||
|
|
np.testing.assert_equal(out, np.ones(2, dtype=np.float32))
|
||
|
|
|
||
|
|
def test_array_pass_through(self):
|
||
|
|
def pyfunc(y):
|
||
|
|
return y
|
||
|
|
|
||
|
|
arr = np.ones(4, dtype=np.float32)
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
expected = cfunc(arr)
|
||
|
|
got = pyfunc(arr)
|
||
|
|
|
||
|
|
np.testing.assert_equal(expected, arr)
|
||
|
|
np.testing.assert_equal(expected, got)
|
||
|
|
self.assertIs(expected, arr)
|
||
|
|
self.assertIs(expected, got)
|
||
|
|
|
||
|
|
def test_array_pass_through_sliced(self):
|
||
|
|
def pyfunc(y):
|
||
|
|
return y[y.size // 2:]
|
||
|
|
|
||
|
|
arr = np.ones(4, dtype=np.float32)
|
||
|
|
|
||
|
|
initrefct = sys.getrefcount(arr)
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
got = cfunc(arr)
|
||
|
|
self.assertEqual(initrefct + 1, sys.getrefcount(arr))
|
||
|
|
expected = pyfunc(arr)
|
||
|
|
self.assertEqual(initrefct + 2, sys.getrefcount(arr))
|
||
|
|
|
||
|
|
np.testing.assert_equal(expected, arr[arr.size // 2])
|
||
|
|
np.testing.assert_equal(expected, got)
|
||
|
|
|
||
|
|
del expected
|
||
|
|
self.assertEqual(initrefct + 1, sys.getrefcount(arr))
|
||
|
|
del got
|
||
|
|
self.assertEqual(initrefct, sys.getrefcount(arr))
|
||
|
|
|
||
|
|
def test_ufunc_with_allocated_output(self):
|
||
|
|
|
||
|
|
def pyfunc(a, b):
|
||
|
|
out = np.empty(a.shape)
|
||
|
|
np.add(a, b, out)
|
||
|
|
return out
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
# 1D case
|
||
|
|
arr_a = np.random.random(10)
|
||
|
|
arr_b = np.random.random(10)
|
||
|
|
|
||
|
|
np.testing.assert_equal(pyfunc(arr_a, arr_b),
|
||
|
|
cfunc(arr_a, arr_b))
|
||
|
|
|
||
|
|
self.assert_array_nrt_refct(cfunc(arr_a, arr_b), 1)
|
||
|
|
|
||
|
|
# 2D case
|
||
|
|
arr_a = np.random.random(10).reshape(2, 5)
|
||
|
|
arr_b = np.random.random(10).reshape(2, 5)
|
||
|
|
|
||
|
|
np.testing.assert_equal(pyfunc(arr_a, arr_b),
|
||
|
|
cfunc(arr_a, arr_b))
|
||
|
|
|
||
|
|
self.assert_array_nrt_refct(cfunc(arr_a, arr_b), 1)
|
||
|
|
|
||
|
|
# 3D case
|
||
|
|
arr_a = np.random.random(70).reshape(2, 5, 7)
|
||
|
|
arr_b = np.random.random(70).reshape(2, 5, 7)
|
||
|
|
|
||
|
|
np.testing.assert_equal(pyfunc(arr_a, arr_b),
|
||
|
|
cfunc(arr_a, arr_b))
|
||
|
|
|
||
|
|
self.assert_array_nrt_refct(cfunc(arr_a, arr_b), 1)
|
||
|
|
|
||
|
|
def test_allocation_mt(self):
|
||
|
|
"""
|
||
|
|
This test exercises the array allocation in multithreaded usecase.
|
||
|
|
This stress the freelist inside NRT.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def pyfunc(inp):
|
||
|
|
out = np.empty(inp.size)
|
||
|
|
|
||
|
|
# Zero fill
|
||
|
|
for i in range(out.size):
|
||
|
|
out[i] = 0
|
||
|
|
|
||
|
|
for i in range(inp[0]):
|
||
|
|
# Allocate inside a loop
|
||
|
|
tmp = np.empty(inp.size)
|
||
|
|
# Write to tmp
|
||
|
|
for j in range(tmp.size):
|
||
|
|
tmp[j] = inp[j]
|
||
|
|
# out = tmp + i
|
||
|
|
for j in range(tmp.size):
|
||
|
|
out[j] += tmp[j] + i
|
||
|
|
|
||
|
|
return out
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
size = 10 # small array size so that the computation is short
|
||
|
|
arr = np.random.randint(1, 10, size)
|
||
|
|
frozen_arr = arr.copy()
|
||
|
|
|
||
|
|
np.testing.assert_equal(pyfunc(arr), cfunc(arr))
|
||
|
|
# Ensure we did not modify the input
|
||
|
|
np.testing.assert_equal(frozen_arr, arr)
|
||
|
|
|
||
|
|
workers = []
|
||
|
|
inputs = []
|
||
|
|
outputs = []
|
||
|
|
|
||
|
|
# Make wrapper to store the output
|
||
|
|
def wrapped(inp, out):
|
||
|
|
out[:] = cfunc(inp)
|
||
|
|
|
||
|
|
# Create a lot of worker threads to create contention
|
||
|
|
for i in range(100):
|
||
|
|
arr = np.random.randint(1, 10, size)
|
||
|
|
out = np.empty_like(arr)
|
||
|
|
thread = threading.Thread(target=wrapped,
|
||
|
|
args=(arr, out),
|
||
|
|
name="worker{0}".format(i))
|
||
|
|
workers.append(thread)
|
||
|
|
inputs.append(arr)
|
||
|
|
outputs.append(out)
|
||
|
|
|
||
|
|
# Launch worker threads
|
||
|
|
for thread in workers:
|
||
|
|
thread.start()
|
||
|
|
|
||
|
|
# Join worker threads
|
||
|
|
for thread in workers:
|
||
|
|
thread.join()
|
||
|
|
|
||
|
|
# Check result
|
||
|
|
for inp, out in zip(inputs, outputs):
|
||
|
|
np.testing.assert_equal(pyfunc(inp), out)
|
||
|
|
|
||
|
|
def test_refct_mt(self):
|
||
|
|
"""
|
||
|
|
This test exercises the refct in multithreaded code
|
||
|
|
"""
|
||
|
|
|
||
|
|
def pyfunc(n, inp):
|
||
|
|
out = np.empty(inp.size)
|
||
|
|
for i in range(out.size):
|
||
|
|
out[i] = inp[i] + 1
|
||
|
|
# Use swap to trigger many refct ops
|
||
|
|
for i in range(n):
|
||
|
|
out, inp = inp, out
|
||
|
|
return out
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
size = 10
|
||
|
|
input = np.arange(size, dtype=float)
|
||
|
|
expected_refct = sys.getrefcount(input)
|
||
|
|
swapct = random.randrange(1000)
|
||
|
|
expected = pyfunc(swapct, input)
|
||
|
|
np.testing.assert_equal(expected, cfunc(swapct, input))
|
||
|
|
# The following checks can discover a reference count error
|
||
|
|
del expected
|
||
|
|
self.assertEqual(expected_refct, sys.getrefcount(input))
|
||
|
|
|
||
|
|
workers = []
|
||
|
|
outputs = []
|
||
|
|
swapcts = []
|
||
|
|
|
||
|
|
# Make wrapper to store the output
|
||
|
|
def wrapped(n, input, out):
|
||
|
|
out[:] = cfunc(n, input)
|
||
|
|
|
||
|
|
# Create worker threads
|
||
|
|
for i in range(100):
|
||
|
|
out = np.empty(size)
|
||
|
|
# All thread shares the same input
|
||
|
|
swapct = random.randrange(1000)
|
||
|
|
thread = threading.Thread(target=wrapped,
|
||
|
|
args=(swapct, input, out),
|
||
|
|
name="worker{0}".format(i))
|
||
|
|
workers.append(thread)
|
||
|
|
outputs.append(out)
|
||
|
|
swapcts.append(swapct)
|
||
|
|
|
||
|
|
# Launch worker threads
|
||
|
|
for thread in workers:
|
||
|
|
thread.start()
|
||
|
|
|
||
|
|
# Join worker threads
|
||
|
|
for thread in workers:
|
||
|
|
thread.join()
|
||
|
|
|
||
|
|
# Check result
|
||
|
|
for swapct, out in zip(swapcts, outputs):
|
||
|
|
np.testing.assert_equal(pyfunc(swapct, input), out)
|
||
|
|
|
||
|
|
del outputs, workers
|
||
|
|
# The following checks can discover a reference count error
|
||
|
|
self.assertEqual(expected_refct, sys.getrefcount(input))
|
||
|
|
|
||
|
|
@skip_if_32bit
|
||
|
|
def test_invalid_size_array(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(x):
|
||
|
|
np.empty(x)
|
||
|
|
|
||
|
|
# Exceptions leak references
|
||
|
|
self.disable_leak_check()
|
||
|
|
|
||
|
|
with self.assertRaises(MemoryError) as raises:
|
||
|
|
foo(types.size_t.maxval // 8 // 2)
|
||
|
|
|
||
|
|
self.assertIn("Allocation failed", str(raises.exception))
|
||
|
|
|
||
|
|
def test_swap(self):
|
||
|
|
|
||
|
|
def pyfunc(x, y, t):
|
||
|
|
"""Swap array x and y for t number of times
|
||
|
|
"""
|
||
|
|
for i in range(t):
|
||
|
|
x, y = y, x
|
||
|
|
|
||
|
|
return x, y
|
||
|
|
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
x = np.random.random(100)
|
||
|
|
y = np.random.random(100)
|
||
|
|
|
||
|
|
t = 100
|
||
|
|
|
||
|
|
initrefct = sys.getrefcount(x), sys.getrefcount(y)
|
||
|
|
expect, got = pyfunc(x, y, t), cfunc(x, y, t)
|
||
|
|
self.assertIsNone(got[0].base)
|
||
|
|
self.assertIsNone(got[1].base)
|
||
|
|
np.testing.assert_equal(expect, got)
|
||
|
|
del expect, got
|
||
|
|
self.assertEqual(initrefct, (sys.getrefcount(x), sys.getrefcount(y)))
|
||
|
|
|
||
|
|
def test_return_tuple_of_array(self):
|
||
|
|
|
||
|
|
def pyfunc(x):
|
||
|
|
y = np.empty(x.size)
|
||
|
|
for i in range(y.size):
|
||
|
|
y[i] = x[i] + 1
|
||
|
|
return x, y
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
x = np.random.random(5)
|
||
|
|
initrefct = sys.getrefcount(x)
|
||
|
|
expected_x, expected_y = pyfunc(x)
|
||
|
|
got_x, got_y = cfunc(x)
|
||
|
|
self.assertIs(x, expected_x)
|
||
|
|
self.assertIs(x, got_x)
|
||
|
|
np.testing.assert_equal(expected_x, got_x)
|
||
|
|
np.testing.assert_equal(expected_y, got_y)
|
||
|
|
del expected_x, got_x
|
||
|
|
self.assertEqual(initrefct, sys.getrefcount(x))
|
||
|
|
|
||
|
|
self.assertEqual(sys.getrefcount(expected_y), sys.getrefcount(got_y))
|
||
|
|
|
||
|
|
def test_return_tuple_of_array_created(self):
|
||
|
|
|
||
|
|
def pyfunc(x):
|
||
|
|
y = np.empty(x.size)
|
||
|
|
for i in range(y.size):
|
||
|
|
y[i] = x[i] + 1
|
||
|
|
out = y, y
|
||
|
|
return out
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
x = np.random.random(5)
|
||
|
|
expected_x, expected_y = pyfunc(x)
|
||
|
|
got_x, got_y = cfunc(x)
|
||
|
|
np.testing.assert_equal(expected_x, got_x)
|
||
|
|
np.testing.assert_equal(expected_y, got_y)
|
||
|
|
# getrefcount owns 1, got_y owns 1
|
||
|
|
self.assertEqual(2, sys.getrefcount(got_y))
|
||
|
|
# getrefcount owns 1, got_y owns 1
|
||
|
|
self.assertEqual(2, sys.getrefcount(got_y))
|
||
|
|
|
||
|
|
def test_issue_with_return_leak(self):
|
||
|
|
"""
|
||
|
|
Dispatcher returns a new reference.
|
||
|
|
It need to workaround it for now.
|
||
|
|
"""
|
||
|
|
@nrtjit
|
||
|
|
def inner(out):
|
||
|
|
return out
|
||
|
|
|
||
|
|
def pyfunc(x):
|
||
|
|
return inner(x)
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
arr = np.arange(10)
|
||
|
|
old_refct = sys.getrefcount(arr)
|
||
|
|
|
||
|
|
self.assertEqual(old_refct, sys.getrefcount(pyfunc(arr)))
|
||
|
|
self.assertEqual(old_refct, sys.getrefcount(cfunc(arr)))
|
||
|
|
self.assertEqual(old_refct, sys.getrefcount(arr))
|
||
|
|
|
||
|
|
|
||
|
|
class ConstructorBaseTest(NrtRefCtTest):
|
||
|
|
|
||
|
|
def check_0d(self, pyfunc):
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
expected = pyfunc()
|
||
|
|
ret = cfunc()
|
||
|
|
self.assert_array_nrt_refct(ret, 1)
|
||
|
|
self.assertEqual(ret.size, expected.size)
|
||
|
|
self.assertEqual(ret.shape, expected.shape)
|
||
|
|
self.assertEqual(ret.dtype, expected.dtype)
|
||
|
|
self.assertEqual(ret.strides, expected.strides)
|
||
|
|
self.check_result_value(ret, expected)
|
||
|
|
# test writability
|
||
|
|
expected = np.empty_like(ret) # np.full_like was not added until Numpy 1.8
|
||
|
|
expected.fill(123)
|
||
|
|
ret.fill(123)
|
||
|
|
np.testing.assert_equal(ret, expected)
|
||
|
|
|
||
|
|
def check_1d(self, pyfunc):
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
n = 3
|
||
|
|
expected = pyfunc(n)
|
||
|
|
ret = cfunc(n)
|
||
|
|
self.assert_array_nrt_refct(ret, 1)
|
||
|
|
self.assertEqual(ret.size, expected.size)
|
||
|
|
self.assertEqual(ret.shape, expected.shape)
|
||
|
|
self.assertEqual(ret.dtype, expected.dtype)
|
||
|
|
self.assertEqual(ret.strides, expected.strides)
|
||
|
|
self.check_result_value(ret, expected)
|
||
|
|
# test writability
|
||
|
|
expected = np.empty_like(ret) # np.full_like was not added until Numpy 1.8
|
||
|
|
expected.fill(123)
|
||
|
|
ret.fill(123)
|
||
|
|
np.testing.assert_equal(ret, expected)
|
||
|
|
# errors
|
||
|
|
with self.assertRaises(ValueError) as cm:
|
||
|
|
cfunc(-1)
|
||
|
|
self.assertEqual(str(cm.exception), "negative dimensions not allowed")
|
||
|
|
|
||
|
|
def check_2d(self, pyfunc):
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
m, n = 2, 3
|
||
|
|
expected = pyfunc(m, n)
|
||
|
|
ret = cfunc(m, n)
|
||
|
|
self.assert_array_nrt_refct(ret, 1)
|
||
|
|
self.assertEqual(ret.size, expected.size)
|
||
|
|
self.assertEqual(ret.shape, expected.shape)
|
||
|
|
self.assertEqual(ret.dtype, expected.dtype)
|
||
|
|
self.assertEqual(ret.strides, expected.strides)
|
||
|
|
self.check_result_value(ret, expected)
|
||
|
|
# test writability
|
||
|
|
expected = np.empty_like(ret) # np.full_like was not added until Numpy 1.8
|
||
|
|
expected.fill(123)
|
||
|
|
ret.fill(123)
|
||
|
|
np.testing.assert_equal(ret, expected)
|
||
|
|
# errors
|
||
|
|
with self.assertRaises(ValueError) as cm:
|
||
|
|
cfunc(2, -1)
|
||
|
|
self.assertEqual(str(cm.exception), "negative dimensions not allowed")
|
||
|
|
|
||
|
|
def check_alloc_size(self, pyfunc):
|
||
|
|
"""Checks that pyfunc will error, not segfaulting due to array size."""
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
with self.assertRaises(ValueError) as e:
|
||
|
|
cfunc()
|
||
|
|
self.assertIn(
|
||
|
|
"array is too big",
|
||
|
|
str(e.exception)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class TestNdZeros(ConstructorBaseTest, TestCase):
|
||
|
|
|
||
|
|
def setUp(self):
|
||
|
|
super(TestNdZeros, self).setUp()
|
||
|
|
self.pyfunc = np.zeros
|
||
|
|
|
||
|
|
def check_result_value(self, ret, expected):
|
||
|
|
np.testing.assert_equal(ret, expected)
|
||
|
|
|
||
|
|
def test_0d(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func():
|
||
|
|
return pyfunc(())
|
||
|
|
self.check_0d(func)
|
||
|
|
|
||
|
|
def test_1d(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(n):
|
||
|
|
return pyfunc(n)
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(n):
|
||
|
|
return pyfunc(n, np.int32)
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype_instance(self):
|
||
|
|
# dtype as numpy dtype, not as scalar class
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
_dtype = np.dtype('int32')
|
||
|
|
def func(n):
|
||
|
|
return pyfunc(n, _dtype)
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype_str(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
_dtype = 'int32'
|
||
|
|
def func(n):
|
||
|
|
return pyfunc(n, _dtype)
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def func(n):
|
||
|
|
return pyfunc(n, 'complex128')
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype_str_alternative_spelling(self):
|
||
|
|
# like test_1d_dtype_str but using the shorthand type spellings
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
_dtype = 'i4'
|
||
|
|
def func(n):
|
||
|
|
return pyfunc(n, _dtype)
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def func(n):
|
||
|
|
return pyfunc(n, 'c8')
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype_str_structured_dtype(self):
|
||
|
|
# test_1d_dtype_str but using a structured dtype
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
_dtype = "i4, (2,3)f8"
|
||
|
|
def func(n):
|
||
|
|
return pyfunc(n, _dtype)
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype_non_const_str(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def func(n, dt):
|
||
|
|
return pyfunc(n, dt)
|
||
|
|
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
func(5, 'int32')
|
||
|
|
|
||
|
|
excstr = str(raises.exception)
|
||
|
|
msg = (f"If np.{self.pyfunc.__name__} dtype is a string it must be a "
|
||
|
|
"string constant.")
|
||
|
|
self.assertIn(msg, excstr)
|
||
|
|
|
||
|
|
def test_1d_dtype_invalid_str(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def func(n):
|
||
|
|
return pyfunc(n, 'ABCDEF')
|
||
|
|
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
func(5)
|
||
|
|
|
||
|
|
excstr = str(raises.exception)
|
||
|
|
self.assertIn("Invalid NumPy dtype specified: 'ABCDEF'", excstr)
|
||
|
|
|
||
|
|
def test_2d(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(m, n):
|
||
|
|
return pyfunc((m, n))
|
||
|
|
self.check_2d(func)
|
||
|
|
|
||
|
|
def test_2d_shape_dtypes(self):
|
||
|
|
# Test for issue #4575
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func1(m, n):
|
||
|
|
return pyfunc((np.int16(m), np.int32(n)))
|
||
|
|
self.check_2d(func1)
|
||
|
|
# Using a 64-bit value checks that 32 bit systems will downcast to intp
|
||
|
|
def func2(m, n):
|
||
|
|
return pyfunc((np.int64(m), np.int8(n)))
|
||
|
|
self.check_2d(func2)
|
||
|
|
# Make sure an error is thrown if we can't downcast safely
|
||
|
|
if config.IS_32BITS:
|
||
|
|
cfunc = nrtjit(lambda m, n: pyfunc((m, n)))
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
cfunc(np.int64(1 << (32 - 1)), 1)
|
||
|
|
|
||
|
|
def test_2d_dtype_kwarg(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(m, n):
|
||
|
|
return pyfunc((m, n), dtype=np.complex64)
|
||
|
|
self.check_2d(func)
|
||
|
|
|
||
|
|
def test_2d_dtype_str_kwarg(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(m, n):
|
||
|
|
return pyfunc((m, n), dtype='complex64')
|
||
|
|
self.check_2d(func)
|
||
|
|
|
||
|
|
def test_2d_dtype_str_kwarg_alternative_spelling(self):
|
||
|
|
# as test_2d_dtype_str_kwarg but with the numpy shorthand type spelling
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(m, n):
|
||
|
|
return pyfunc((m, n), dtype='c8')
|
||
|
|
self.check_2d(func)
|
||
|
|
|
||
|
|
def test_alloc_size(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
width = types.intp.bitwidth
|
||
|
|
def gen_func(shape, dtype):
|
||
|
|
return lambda : pyfunc(shape, dtype)
|
||
|
|
# Under these values numba will segfault, but that's another issue
|
||
|
|
self.check_alloc_size(gen_func(1 << width - 2, np.intp))
|
||
|
|
self.check_alloc_size(gen_func((1 << width - 8, 64), np.intp))
|
||
|
|
|
||
|
|
|
||
|
|
class TestNdOnes(TestNdZeros):
|
||
|
|
|
||
|
|
def setUp(self):
|
||
|
|
super(TestNdOnes, self).setUp()
|
||
|
|
self.pyfunc = np.ones
|
||
|
|
|
||
|
|
@unittest.expectedFailure
|
||
|
|
def test_1d_dtype_str_structured_dtype(self):
|
||
|
|
super().test_1d_dtype_str_structured_dtype()
|
||
|
|
|
||
|
|
|
||
|
|
class TestNdFull(ConstructorBaseTest, TestCase):
|
||
|
|
|
||
|
|
def check_result_value(self, ret, expected):
|
||
|
|
np.testing.assert_equal(ret, expected)
|
||
|
|
|
||
|
|
def test_0d(self):
|
||
|
|
def func():
|
||
|
|
return np.full((), 4.5)
|
||
|
|
self.check_0d(func)
|
||
|
|
|
||
|
|
def test_1d(self):
|
||
|
|
def func(n):
|
||
|
|
return np.full(n, 4.5)
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype(self):
|
||
|
|
def func(n):
|
||
|
|
return np.full(n, 4.5, np.bool_)
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype_instance(self):
|
||
|
|
dtype = np.dtype('bool')
|
||
|
|
def func(n):
|
||
|
|
return np.full(n, 4.5, dtype)
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype_str(self):
|
||
|
|
def func(n):
|
||
|
|
return np.full(n, 4.5, 'bool_')
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype_str_alternative_spelling(self):
|
||
|
|
# like test_1d_dtype_str but using the shorthand type spelling
|
||
|
|
def func(n):
|
||
|
|
return np.full(n, 4.5, '?')
|
||
|
|
self.check_1d(func)
|
||
|
|
|
||
|
|
def test_1d_dtype_non_const_str(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def func(n, fv, dt):
|
||
|
|
return np.full(n, fv, dt)
|
||
|
|
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
func((5,), 4.5, 'int32')
|
||
|
|
|
||
|
|
excstr = str(raises.exception)
|
||
|
|
msg = ("If np.full dtype is a string it must be a "
|
||
|
|
"string constant.")
|
||
|
|
self.assertIn(msg, excstr)
|
||
|
|
|
||
|
|
def test_1d_dtype_invalid_str(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def func(n, fv):
|
||
|
|
return np.full(n, fv, 'ABCDEF')
|
||
|
|
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
func((5,), 4.5)
|
||
|
|
|
||
|
|
excstr = str(raises.exception)
|
||
|
|
self.assertIn("Invalid NumPy dtype specified: 'ABCDEF'", excstr)
|
||
|
|
|
||
|
|
def test_2d(self):
|
||
|
|
def func(m, n):
|
||
|
|
return np.full((m, n), 4.5)
|
||
|
|
self.check_2d(func)
|
||
|
|
|
||
|
|
def test_2d_dtype_kwarg(self):
|
||
|
|
def func(m, n):
|
||
|
|
return np.full((m, n), 1 + 4.5j, dtype=np.complex64)
|
||
|
|
self.check_2d(func)
|
||
|
|
|
||
|
|
def test_2d_dtype_from_type(self):
|
||
|
|
# tests issue #2862
|
||
|
|
def func(m, n):
|
||
|
|
return np.full((m, n), np.int32(1))
|
||
|
|
self.check_2d(func)
|
||
|
|
|
||
|
|
# Complex uses `.real`, imaginary part dropped
|
||
|
|
def func(m, n):
|
||
|
|
return np.full((m, n), np.complex128(1))
|
||
|
|
self.check_2d(func)
|
||
|
|
|
||
|
|
# and that if a dtype is specified, this influences the return type
|
||
|
|
def func(m, n):
|
||
|
|
return np.full((m, n), 1, dtype=np.int8)
|
||
|
|
self.check_2d(func)
|
||
|
|
|
||
|
|
def test_2d_shape_dtypes(self):
|
||
|
|
# Test for issue #4575
|
||
|
|
def func1(m, n):
|
||
|
|
return np.full((np.int16(m), np.int32(n)), 4.5)
|
||
|
|
self.check_2d(func1)
|
||
|
|
# Using a 64-bit value checks that 32 bit systems will downcast to intp
|
||
|
|
def func2(m, n):
|
||
|
|
return np.full((np.int64(m), np.int8(n)), 4.5)
|
||
|
|
self.check_2d(func2)
|
||
|
|
# Make sure an error is thrown if we can't downcast safely
|
||
|
|
if config.IS_32BITS:
|
||
|
|
cfunc = nrtjit(lambda m, n: np.full((m, n), 4.5))
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
cfunc(np.int64(1 << (32 - 1)), 1)
|
||
|
|
|
||
|
|
def test_alloc_size(self):
|
||
|
|
width = types.intp.bitwidth
|
||
|
|
def gen_func(shape, value):
|
||
|
|
return lambda : np.full(shape, value)
|
||
|
|
# Under these values numba will segfault, but that's another issue
|
||
|
|
self.check_alloc_size(gen_func(1 << width - 2, 1))
|
||
|
|
self.check_alloc_size(gen_func((1 << width - 8, 64), 1))
|
||
|
|
|
||
|
|
|
||
|
|
class ConstructorLikeBaseTest(object):
|
||
|
|
|
||
|
|
def mutate_array(self, arr):
|
||
|
|
try:
|
||
|
|
arr.fill(42)
|
||
|
|
except (TypeError, ValueError):
|
||
|
|
# Try something else (e.g. Numpy 1.6 with structured dtypes)
|
||
|
|
fill_value = b'x' * arr.dtype.itemsize
|
||
|
|
arr.fill(fill_value)
|
||
|
|
|
||
|
|
def check_like(self, pyfunc, dtype):
|
||
|
|
def check_arr(arr):
|
||
|
|
expected = pyfunc(arr)
|
||
|
|
ret = cfunc(arr)
|
||
|
|
self.assertEqual(ret.size, expected.size)
|
||
|
|
self.assertEqual(ret.dtype, expected.dtype)
|
||
|
|
self.assertStridesEqual(ret, expected)
|
||
|
|
self.check_result_value(ret, expected)
|
||
|
|
# test writability
|
||
|
|
self.mutate_array(ret)
|
||
|
|
self.mutate_array(expected)
|
||
|
|
np.testing.assert_equal(ret, expected)
|
||
|
|
|
||
|
|
orig = np.linspace(0, 5, 6).astype(dtype)
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
for shape in (6, (2, 3), (1, 2, 3), (3, 1, 2), ()):
|
||
|
|
if shape == ():
|
||
|
|
arr = orig[-1:].reshape(())
|
||
|
|
else:
|
||
|
|
arr = orig.reshape(shape)
|
||
|
|
check_arr(arr)
|
||
|
|
# Non-contiguous array
|
||
|
|
if arr.ndim > 0:
|
||
|
|
check_arr(arr[::2])
|
||
|
|
# Check new array doesn't inherit readonly flag
|
||
|
|
arr.flags['WRITEABLE'] = False
|
||
|
|
# verify read-only
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
arr[0] = 1
|
||
|
|
check_arr(arr)
|
||
|
|
|
||
|
|
# Scalar argument => should produce a 0-d array
|
||
|
|
check_arr(orig[0])
|
||
|
|
|
||
|
|
|
||
|
|
class TestNdEmptyLike(ConstructorLikeBaseTest, TestCase):
|
||
|
|
|
||
|
|
def setUp(self):
|
||
|
|
super(TestNdEmptyLike, self).setUp()
|
||
|
|
self.pyfunc = np.empty_like
|
||
|
|
|
||
|
|
def check_result_value(self, ret, expected):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def test_like(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(arr):
|
||
|
|
return pyfunc(arr)
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_structured(self):
|
||
|
|
dtype = np.dtype([('a', np.int16), ('b', np.float32)])
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(arr):
|
||
|
|
return pyfunc(arr)
|
||
|
|
self.check_like(func, dtype)
|
||
|
|
|
||
|
|
def test_like_dtype(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(arr):
|
||
|
|
return pyfunc(arr, np.int32)
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_instance(self):
|
||
|
|
dtype = np.dtype('int32')
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(arr):
|
||
|
|
return pyfunc(arr, dtype)
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_structured(self):
|
||
|
|
dtype = np.dtype([('a', np.int16), ('b', np.float32)])
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(arr):
|
||
|
|
return pyfunc(arr, dtype)
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_kwarg(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(arr):
|
||
|
|
return pyfunc(arr, dtype=np.int32)
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_str_kwarg(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(arr):
|
||
|
|
return pyfunc(arr, dtype='int32')
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_str_kwarg_alternative_spelling(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
def func(arr):
|
||
|
|
return pyfunc(arr, dtype='i4')
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_non_const_str(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def func(n, dt):
|
||
|
|
return pyfunc(n, dt)
|
||
|
|
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
func(np.ones(4), 'int32')
|
||
|
|
|
||
|
|
excstr = str(raises.exception)
|
||
|
|
msg = (f"If np.{self.pyfunc.__name__} dtype is a string it must be a "
|
||
|
|
"string constant.")
|
||
|
|
self.assertIn(msg, excstr)
|
||
|
|
self.assertIn(
|
||
|
|
'{}(array(float64, 1d, C), unicode_type)'.format(pyfunc.__name__),
|
||
|
|
excstr)
|
||
|
|
|
||
|
|
def test_like_dtype_invalid_str(self):
|
||
|
|
pyfunc = self.pyfunc
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def func(n):
|
||
|
|
return pyfunc(n, 'ABCDEF')
|
||
|
|
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
func(np.ones(4))
|
||
|
|
|
||
|
|
excstr = str(raises.exception)
|
||
|
|
self.assertIn("Invalid NumPy dtype specified: 'ABCDEF'", excstr)
|
||
|
|
|
||
|
|
|
||
|
|
class TestNdZerosLike(TestNdEmptyLike):
|
||
|
|
|
||
|
|
def setUp(self):
|
||
|
|
super(TestNdZerosLike, self).setUp()
|
||
|
|
self.pyfunc = np.zeros_like
|
||
|
|
|
||
|
|
def check_result_value(self, ret, expected):
|
||
|
|
np.testing.assert_equal(ret, expected)
|
||
|
|
|
||
|
|
def test_like_structured(self):
|
||
|
|
super(TestNdZerosLike, self).test_like_structured()
|
||
|
|
|
||
|
|
def test_like_dtype_structured(self):
|
||
|
|
super(TestNdZerosLike, self).test_like_dtype_structured()
|
||
|
|
|
||
|
|
|
||
|
|
class TestNdOnesLike(TestNdZerosLike):
|
||
|
|
|
||
|
|
def setUp(self):
|
||
|
|
super(TestNdOnesLike, self).setUp()
|
||
|
|
self.pyfunc = np.ones_like
|
||
|
|
self.expected_value = 1
|
||
|
|
|
||
|
|
# Not supported yet.
|
||
|
|
|
||
|
|
@unittest.expectedFailure
|
||
|
|
def test_like_structured(self):
|
||
|
|
super(TestNdOnesLike, self).test_like_structured()
|
||
|
|
|
||
|
|
@unittest.expectedFailure
|
||
|
|
def test_like_dtype_structured(self):
|
||
|
|
super(TestNdOnesLike, self).test_like_dtype_structured()
|
||
|
|
|
||
|
|
|
||
|
|
class TestNdFullLike(ConstructorLikeBaseTest, TestCase):
|
||
|
|
|
||
|
|
def check_result_value(self, ret, expected):
|
||
|
|
np.testing.assert_equal(ret, expected)
|
||
|
|
|
||
|
|
def test_like(self):
|
||
|
|
def func(arr):
|
||
|
|
return np.full_like(arr, 3.5)
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
# Not supported yet.
|
||
|
|
@unittest.expectedFailure
|
||
|
|
def test_like_structured(self):
|
||
|
|
dtype = np.dtype([('a', np.int16), ('b', np.float32)])
|
||
|
|
def func(arr):
|
||
|
|
return np.full_like(arr, 4.5)
|
||
|
|
self.check_like(func, dtype)
|
||
|
|
|
||
|
|
def test_like_dtype(self):
|
||
|
|
def func(arr):
|
||
|
|
return np.full_like(arr, 4.5, np.bool_)
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_instance(self):
|
||
|
|
dtype = np.dtype('bool')
|
||
|
|
def func(arr):
|
||
|
|
return np.full_like(arr, 4.5, dtype)
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_kwarg(self):
|
||
|
|
def func(arr):
|
||
|
|
return np.full_like(arr, 4.5, dtype=np.bool_)
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_str_kwarg(self):
|
||
|
|
def func(arr):
|
||
|
|
return np.full_like(arr, 4.5, 'bool_')
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_str_kwarg_alternative_spelling(self):
|
||
|
|
def func(arr):
|
||
|
|
return np.full_like(arr, 4.5, dtype='?')
|
||
|
|
self.check_like(func, np.float64)
|
||
|
|
|
||
|
|
def test_like_dtype_non_const_str_kwarg(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def func(arr, fv, dt):
|
||
|
|
return np.full_like(arr, fv, dt)
|
||
|
|
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
func(np.ones(3,), 4.5, 'int32')
|
||
|
|
|
||
|
|
excstr = str(raises.exception)
|
||
|
|
msg = ("If np.full_like dtype is a string it must be a "
|
||
|
|
"string constant.")
|
||
|
|
self.assertIn(msg, excstr)
|
||
|
|
|
||
|
|
def test_like_dtype_invalid_str(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def func(arr, fv):
|
||
|
|
return np.full_like(arr, fv, "ABCDEF")
|
||
|
|
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
func(np.ones(4), 3.4)
|
||
|
|
|
||
|
|
excstr = str(raises.exception)
|
||
|
|
self.assertIn("Invalid NumPy dtype specified: 'ABCDEF'", excstr)
|
||
|
|
|
||
|
|
|
||
|
|
class TestNdIdentity(BaseTest):
|
||
|
|
|
||
|
|
def check_identity(self, pyfunc):
|
||
|
|
self.check_outputs(pyfunc, [(3,)])
|
||
|
|
|
||
|
|
def test_identity(self):
|
||
|
|
def func(n):
|
||
|
|
return np.identity(n)
|
||
|
|
self.check_identity(func)
|
||
|
|
|
||
|
|
def test_identity_dtype(self):
|
||
|
|
for dtype in (np.complex64, np.int16, np.bool_, np.dtype('bool'),
|
||
|
|
'bool_'):
|
||
|
|
def func(n):
|
||
|
|
return np.identity(n, dtype)
|
||
|
|
self.check_identity(func)
|
||
|
|
|
||
|
|
def test_like_dtype_non_const_str_kwarg(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def func(n, dt):
|
||
|
|
return np.identity(n, dt)
|
||
|
|
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
func(4, 'int32')
|
||
|
|
|
||
|
|
excstr = str(raises.exception)
|
||
|
|
msg = ("If np.identity dtype is a string it must be a "
|
||
|
|
"string constant.")
|
||
|
|
self.assertIn(msg, excstr)
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
class TestNdEye(BaseTest):
|
||
|
|
|
||
|
|
def test_eye_n(self):
|
||
|
|
def func(n):
|
||
|
|
return np.eye(n)
|
||
|
|
self.check_outputs(func, [(1,), (3,)])
|
||
|
|
|
||
|
|
def test_eye_n_dtype(self):
|
||
|
|
# check None option, dtype class, instance of dtype class
|
||
|
|
for dt in (None, np.complex128, np.complex64(1)):
|
||
|
|
def func(n, dtype=dt):
|
||
|
|
return np.eye(n, dtype=dtype)
|
||
|
|
self.check_outputs(func, [(1,), (3,)])
|
||
|
|
|
||
|
|
def test_eye_n_m(self):
|
||
|
|
def func(n, m):
|
||
|
|
return np.eye(n, m)
|
||
|
|
self.check_outputs(func, [(1, 2), (3, 2), (0, 3)])
|
||
|
|
|
||
|
|
def check_eye_n_m_k(self, func):
|
||
|
|
self.check_outputs(func, [(1, 2, 0),
|
||
|
|
(3, 4, 1),
|
||
|
|
(3, 4, -1),
|
||
|
|
(4, 3, -2),
|
||
|
|
(4, 3, -5),
|
||
|
|
(4, 3, 5)])
|
||
|
|
|
||
|
|
def test_eye_n_m_k(self):
|
||
|
|
def func(n, m, k):
|
||
|
|
return np.eye(n, m, k)
|
||
|
|
self.check_eye_n_m_k(func)
|
||
|
|
|
||
|
|
def test_eye_n_m_k_dtype(self):
|
||
|
|
def func(n, m, k):
|
||
|
|
return np.eye(N=n, M=m, k=k, dtype=np.int16)
|
||
|
|
self.check_eye_n_m_k(func)
|
||
|
|
|
||
|
|
def test_eye_n_m_k_dtype_instance(self):
|
||
|
|
dtype = np.dtype('int16')
|
||
|
|
def func(n, m, k):
|
||
|
|
return np.eye(N=n, M=m, k=k, dtype=dtype)
|
||
|
|
self.check_eye_n_m_k(func)
|
||
|
|
|
||
|
|
|
||
|
|
class TestNdDiag(TestCase):
|
||
|
|
|
||
|
|
def setUp(self):
|
||
|
|
v = np.array([1, 2, 3])
|
||
|
|
hv = np.array([[1, 2, 3]])
|
||
|
|
vv = np.transpose(hv)
|
||
|
|
self.vectors = [v, hv, vv]
|
||
|
|
a3x4 = np.arange(12).reshape(3, 4)
|
||
|
|
a4x3 = np.arange(12).reshape(4, 3)
|
||
|
|
self.matricies = [a3x4, a4x3]
|
||
|
|
def func(q):
|
||
|
|
return np.diag(q)
|
||
|
|
self.py = func
|
||
|
|
self.jit = nrtjit(func)
|
||
|
|
|
||
|
|
def func_kwarg(q, k=0):
|
||
|
|
return np.diag(q, k=k)
|
||
|
|
self.py_kw = func_kwarg
|
||
|
|
self.jit_kw = nrtjit(func_kwarg)
|
||
|
|
|
||
|
|
def check_diag(self, pyfunc, nrtfunc, *args, **kwargs):
|
||
|
|
expected = pyfunc(*args, **kwargs)
|
||
|
|
computed = nrtfunc(*args, **kwargs)
|
||
|
|
self.assertEqual(computed.size, expected.size)
|
||
|
|
self.assertEqual(computed.dtype, expected.dtype)
|
||
|
|
# NOTE: stride not tested as np returns a RO view, nb returns new data
|
||
|
|
np.testing.assert_equal(expected, computed)
|
||
|
|
|
||
|
|
# create a diag matrix from a vector
|
||
|
|
def test_diag_vect_create(self):
|
||
|
|
for d in self.vectors:
|
||
|
|
self.check_diag(self.py, self.jit, d)
|
||
|
|
|
||
|
|
# create a diag matrix from a vector at a given offset
|
||
|
|
def test_diag_vect_create_kwarg(self):
|
||
|
|
for k in range(-10, 10):
|
||
|
|
for d in self.vectors:
|
||
|
|
self.check_diag(self.py_kw, self.jit_kw, d, k=k)
|
||
|
|
|
||
|
|
# extract the diagonal
|
||
|
|
def test_diag_extract(self):
|
||
|
|
for d in self.matricies:
|
||
|
|
self.check_diag(self.py, self.jit, d)
|
||
|
|
|
||
|
|
# extract a diagonal at a given offset
|
||
|
|
def test_diag_extract_kwarg(self):
|
||
|
|
for k in range(-4, 4):
|
||
|
|
for d in self.matricies:
|
||
|
|
self.check_diag(self.py_kw, self.jit_kw, d, k=k)
|
||
|
|
|
||
|
|
# check error handling
|
||
|
|
def test_error_handling(self):
|
||
|
|
d = np.array([[[1.]]])
|
||
|
|
cfunc = nrtjit(self.py)
|
||
|
|
|
||
|
|
# missing arg
|
||
|
|
with self.assertRaises(TypeError):
|
||
|
|
cfunc()
|
||
|
|
|
||
|
|
# > 2d
|
||
|
|
with self.assertRaises(TypingError):
|
||
|
|
cfunc(d)
|
||
|
|
with self.assertRaises(TypingError):
|
||
|
|
dfunc = nrtjit(self.py_kw)
|
||
|
|
dfunc(d, k=3)
|
||
|
|
|
||
|
|
def test_bad_shape(self):
|
||
|
|
cfunc = nrtjit(self.py)
|
||
|
|
msg = '.*The argument "v" must be array-like.*'
|
||
|
|
with self.assertRaisesRegex(TypingError, msg) as raises:
|
||
|
|
cfunc(None)
|
||
|
|
|
||
|
|
class TestLinspace(BaseTest):
|
||
|
|
|
||
|
|
def test_linspace_2(self):
|
||
|
|
def pyfunc(n, m):
|
||
|
|
return np.linspace(n, m)
|
||
|
|
self.check_outputs(pyfunc,
|
||
|
|
[(0, 4), (1, 100), (-3.5, 2.5), (-3j, 2+3j),
|
||
|
|
(2, 1), (1+0.5j, 1.5j)])
|
||
|
|
|
||
|
|
def test_linspace_3(self):
|
||
|
|
def pyfunc(n, m, p):
|
||
|
|
return np.linspace(n, m, p)
|
||
|
|
self.check_outputs(pyfunc,
|
||
|
|
[(0, 4, 9), (1, 4, 3), (-3.5, 2.5, 8),
|
||
|
|
(-3j, 2+3j, 7), (2, 1, 0),
|
||
|
|
(1+0.5j, 1.5j, 5), (1, 1e100, 1)])
|
||
|
|
|
||
|
|
def test_linspace_accuracy(self):
|
||
|
|
# Checking linspace reasonably replicates NumPy's algorithm
|
||
|
|
# see https://github.com/numba/numba/issues/6768
|
||
|
|
@nrtjit
|
||
|
|
def foo(n, m, p):
|
||
|
|
return np.linspace(n, m, p)
|
||
|
|
|
||
|
|
n, m, p = 0.0, 1.0, 100
|
||
|
|
self.assertPreciseEqual(foo(n, m, p), foo.py_func(n, m, p))
|
||
|
|
|
||
|
|
|
||
|
|
class TestNpyEmptyKeyword(TestCase):
|
||
|
|
def _test_with_dtype_kw(self, dtype):
|
||
|
|
def pyfunc(shape):
|
||
|
|
return np.empty(shape, dtype=dtype)
|
||
|
|
|
||
|
|
shapes = [1, 5, 9]
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
for s in shapes:
|
||
|
|
expected = pyfunc(s)
|
||
|
|
got = cfunc(s)
|
||
|
|
self.assertEqual(expected.dtype, got.dtype)
|
||
|
|
self.assertEqual(expected.shape, got.shape)
|
||
|
|
|
||
|
|
def test_with_dtype_kws(self):
|
||
|
|
for dtype in [np.int32, np.float32, np.complex64, np.dtype('complex64')]:
|
||
|
|
self._test_with_dtype_kw(dtype)
|
||
|
|
|
||
|
|
def _test_with_shape_and_dtype_kw(self, dtype):
|
||
|
|
def pyfunc(shape):
|
||
|
|
return np.empty(shape=shape, dtype=dtype)
|
||
|
|
|
||
|
|
shapes = [1, 5, 9]
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
for s in shapes:
|
||
|
|
expected = pyfunc(s)
|
||
|
|
got = cfunc(s)
|
||
|
|
self.assertEqual(expected.dtype, got.dtype)
|
||
|
|
self.assertEqual(expected.shape, got.shape)
|
||
|
|
|
||
|
|
def test_with_shape_and_dtype_kws(self):
|
||
|
|
for dtype in [np.int32, np.float32, np.complex64, np.dtype('complex64')]:
|
||
|
|
self._test_with_shape_and_dtype_kw(dtype)
|
||
|
|
|
||
|
|
def test_empty_no_args(self):
|
||
|
|
|
||
|
|
def pyfunc():
|
||
|
|
return np.empty()
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
# Trigger the compilation
|
||
|
|
# That will cause a TypingError due to missing shape argument
|
||
|
|
with self.assertRaises(TypingError):
|
||
|
|
cfunc()
|
||
|
|
|
||
|
|
|
||
|
|
class TestNpArray(MemoryLeakMixin, BaseTest):
|
||
|
|
|
||
|
|
def test_0d(self):
|
||
|
|
def pyfunc(arg):
|
||
|
|
return np.array(arg)
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
got = cfunc(42)
|
||
|
|
self.assertPreciseEqual(got, np.array(42, dtype=np.intp))
|
||
|
|
got = cfunc(2.5)
|
||
|
|
self.assertPreciseEqual(got, np.array(2.5))
|
||
|
|
|
||
|
|
def test_0d_with_dtype(self):
|
||
|
|
def pyfunc(arg):
|
||
|
|
return np.array(arg, dtype=np.int16)
|
||
|
|
|
||
|
|
self.check_outputs(pyfunc, [(42,), (3.5,)])
|
||
|
|
|
||
|
|
def test_1d(self):
|
||
|
|
def pyfunc(arg):
|
||
|
|
return np.array(arg)
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
# A list
|
||
|
|
got = cfunc([2, 3, 42])
|
||
|
|
self.assertPreciseEqual(got, np.intp([2, 3, 42]))
|
||
|
|
# A heterogeneous tuple
|
||
|
|
got = cfunc((1.0, 2.5j, 42))
|
||
|
|
self.assertPreciseEqual(got, np.array([1.0, 2.5j, 42]))
|
||
|
|
# An empty tuple
|
||
|
|
got = cfunc(())
|
||
|
|
self.assertPreciseEqual(got, np.float64(()))
|
||
|
|
|
||
|
|
def test_1d_with_dtype(self):
|
||
|
|
def pyfunc(arg):
|
||
|
|
return np.array(arg, dtype=np.float32)
|
||
|
|
|
||
|
|
self.check_outputs(pyfunc,
|
||
|
|
[([2, 42],),
|
||
|
|
([3.5, 1.0],),
|
||
|
|
((1, 3.5, 42),),
|
||
|
|
((),),
|
||
|
|
])
|
||
|
|
|
||
|
|
def test_1d_with_str_dtype(self):
|
||
|
|
def pyfunc(arg):
|
||
|
|
return np.array(arg, dtype='float32')
|
||
|
|
|
||
|
|
self.check_outputs(pyfunc,
|
||
|
|
[([2, 42],),
|
||
|
|
([3.5, 1.0],),
|
||
|
|
((1, 3.5, 42),),
|
||
|
|
((),),
|
||
|
|
])
|
||
|
|
|
||
|
|
def test_1d_with_non_const_str_dtype(self):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def func(arg, dt):
|
||
|
|
return np.array(arg, dtype=dt)
|
||
|
|
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
func((5, 3), 'int32')
|
||
|
|
|
||
|
|
excstr = str(raises.exception)
|
||
|
|
msg = (f"If np.array dtype is a string it must be a "
|
||
|
|
"string constant.")
|
||
|
|
self.assertIn(msg, excstr)
|
||
|
|
|
||
|
|
def test_2d(self):
|
||
|
|
def pyfunc(arg):
|
||
|
|
return np.array(arg)
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
# A list of tuples
|
||
|
|
got = cfunc([(1, 2), (3, 4)])
|
||
|
|
self.assertPreciseEqual(got, np.intp([[1, 2], [3, 4]]))
|
||
|
|
got = cfunc([(1, 2.5), (3, 4.5)])
|
||
|
|
self.assertPreciseEqual(got, np.float64([[1, 2.5], [3, 4.5]]))
|
||
|
|
# A tuple of lists
|
||
|
|
got = cfunc(([1, 2], [3, 4]))
|
||
|
|
self.assertPreciseEqual(got, np.intp([[1, 2], [3, 4]]))
|
||
|
|
got = cfunc(([1, 2], [3.5, 4.5]))
|
||
|
|
self.assertPreciseEqual(got, np.float64([[1, 2], [3.5, 4.5]]))
|
||
|
|
# A tuple of tuples
|
||
|
|
got = cfunc(((1.5, 2), (3.5, 4.5)))
|
||
|
|
self.assertPreciseEqual(got, np.float64([[1.5, 2], [3.5, 4.5]]))
|
||
|
|
got = cfunc(((), ()))
|
||
|
|
self.assertPreciseEqual(got, np.float64(((), ())))
|
||
|
|
|
||
|
|
def test_2d_with_dtype(self):
|
||
|
|
def pyfunc(arg):
|
||
|
|
return np.array(arg, dtype=np.int32)
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
got = cfunc([(1, 2.5), (3, 4.5)])
|
||
|
|
self.assertPreciseEqual(got, np.int32([[1, 2], [3, 4]]))
|
||
|
|
|
||
|
|
def test_raises(self):
|
||
|
|
|
||
|
|
def pyfunc(arg):
|
||
|
|
return np.array(arg)
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
@contextlib.contextmanager
|
||
|
|
def check_raises(msg):
|
||
|
|
with self.assertRaises(TypingError) as raises:
|
||
|
|
yield
|
||
|
|
self.assertIn(msg, str(raises.exception))
|
||
|
|
|
||
|
|
with check_raises(('array(float64, 1d, C) not allowed in a '
|
||
|
|
'homogeneous sequence')):
|
||
|
|
cfunc(np.array([1.]))
|
||
|
|
|
||
|
|
with check_raises(('type Tuple(int64, reflected list(int64)<iv=None>) '
|
||
|
|
'does not have a regular shape')):
|
||
|
|
cfunc((np.int64(1), [np.int64(2)]))
|
||
|
|
|
||
|
|
with check_raises(
|
||
|
|
"cannot convert Tuple(int64, Record(a[type=int32;offset=0],"
|
||
|
|
"b[type=float32;offset=4];8;False)) to a homogeneous type",
|
||
|
|
):
|
||
|
|
st = np.dtype([('a', 'i4'), ('b', 'f4')])
|
||
|
|
val = np.zeros(1, dtype=st)[0]
|
||
|
|
cfunc(((1, 2), (np.int64(1), val)))
|
||
|
|
|
||
|
|
def test_bad_array(self):
|
||
|
|
@njit
|
||
|
|
def func(obj):
|
||
|
|
return np.array(obj)
|
||
|
|
|
||
|
|
msg = '.*The argument "object" must be array-like.*'
|
||
|
|
with self.assertRaisesRegex(TypingError, msg) as raises:
|
||
|
|
func(None)
|
||
|
|
|
||
|
|
def test_bad_dtype(self):
|
||
|
|
@njit
|
||
|
|
def func(obj, dt):
|
||
|
|
return np.array(obj, dt)
|
||
|
|
|
||
|
|
msg = '.*The argument "dtype" must be a data-type if it is provided.*'
|
||
|
|
with self.assertRaisesRegex(TypingError, msg) as raises:
|
||
|
|
func(5, 4)
|
||
|
|
|
||
|
|
|
||
|
|
class TestNpConcatenate(MemoryLeakMixin, TestCase):
|
||
|
|
"""
|
||
|
|
Tests for np.concatenate().
|
||
|
|
"""
|
||
|
|
|
||
|
|
def _3d_arrays(self):
|
||
|
|
a = np.arange(24).reshape((4, 3, 2))
|
||
|
|
b = a + 10
|
||
|
|
c = (b + 10).copy(order='F')
|
||
|
|
d = (c + 10)[::-1]
|
||
|
|
e = (d + 10)[...,::-1]
|
||
|
|
return a, b, c, d, e
|
||
|
|
|
||
|
|
@contextlib.contextmanager
|
||
|
|
def assert_invalid_sizes_over_dim(self, axis):
|
||
|
|
with self.assertRaises(ValueError) as raises:
|
||
|
|
yield
|
||
|
|
self.assertIn("input sizes over dimension %d do not match" % axis,
|
||
|
|
str(raises.exception))
|
||
|
|
|
||
|
|
def test_3d(self):
|
||
|
|
pyfunc = np_concatenate2
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
def check(a, b, c, axis):
|
||
|
|
for ax in (axis, -3 + axis):
|
||
|
|
expected = pyfunc(a, b, c, axis=ax)
|
||
|
|
got = cfunc(a, b, c, axis=ax)
|
||
|
|
self.assertPreciseEqual(got, expected)
|
||
|
|
|
||
|
|
def check_all_axes(a, b, c):
|
||
|
|
for axis in range(3):
|
||
|
|
check(a, b, c, axis)
|
||
|
|
|
||
|
|
a, b, c, d, e = self._3d_arrays()
|
||
|
|
|
||
|
|
# Inputs with equal sizes
|
||
|
|
# C, C, C
|
||
|
|
check_all_axes(a, b, b)
|
||
|
|
# C, C, F
|
||
|
|
check_all_axes(a, b, c)
|
||
|
|
# F, F, F
|
||
|
|
check_all_axes(a.T, b.T, a.T)
|
||
|
|
# F, F, C
|
||
|
|
check_all_axes(a.T, b.T, c.T)
|
||
|
|
# F, F, A
|
||
|
|
check_all_axes(a.T, b.T, d.T)
|
||
|
|
# A, A, A
|
||
|
|
# (note Numpy may select the layout differently for other inputs)
|
||
|
|
check_all_axes(d.T, e.T, d.T)
|
||
|
|
|
||
|
|
# Inputs with compatible sizes
|
||
|
|
check(a[1:], b, c[::-1], axis=0)
|
||
|
|
check(a, b[:,1:], c, axis=1)
|
||
|
|
check(a, b, c[:,:,1:], axis=2)
|
||
|
|
|
||
|
|
# Different but compatible dtypes
|
||
|
|
check_all_axes(a, b.astype(np.float64), b)
|
||
|
|
|
||
|
|
# Exceptions leak references
|
||
|
|
self.disable_leak_check()
|
||
|
|
|
||
|
|
# Incompatible sizes
|
||
|
|
for axis in (1, 2, -2, -1):
|
||
|
|
with self.assert_invalid_sizes_over_dim(0):
|
||
|
|
cfunc(a[1:], b, b, axis)
|
||
|
|
for axis in (0, 2, -3, -1):
|
||
|
|
with self.assert_invalid_sizes_over_dim(1):
|
||
|
|
cfunc(a, b[:,1:], b, axis)
|
||
|
|
|
||
|
|
def test_3d_no_axis(self):
|
||
|
|
pyfunc = np_concatenate1
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
def check(a, b, c):
|
||
|
|
expected = pyfunc(a, b, c)
|
||
|
|
got = cfunc(a, b, c)
|
||
|
|
self.assertPreciseEqual(got, expected)
|
||
|
|
|
||
|
|
a, b, c, d, e = self._3d_arrays()
|
||
|
|
|
||
|
|
# Inputs with equal sizes
|
||
|
|
# C, C, C
|
||
|
|
check(a, b, b)
|
||
|
|
# C, C, F
|
||
|
|
check(a, b, c)
|
||
|
|
# F, F, F
|
||
|
|
check(a.T, b.T, a.T)
|
||
|
|
# F, F, C
|
||
|
|
check(a.T, b.T, c.T)
|
||
|
|
# F, F, A
|
||
|
|
check(a.T, b.T, d.T)
|
||
|
|
# A, A, A
|
||
|
|
# (note Numpy may select the layout differently for other inputs)
|
||
|
|
check(d.T, e.T, d.T)
|
||
|
|
|
||
|
|
# Inputs with compatible sizes
|
||
|
|
check(a[1:], b, c[::-1])
|
||
|
|
|
||
|
|
# Exceptions leak references
|
||
|
|
self.disable_leak_check()
|
||
|
|
|
||
|
|
# Incompatible sizes
|
||
|
|
with self.assert_invalid_sizes_over_dim(1):
|
||
|
|
cfunc(a, b[:,1:], b)
|
||
|
|
|
||
|
|
def test_typing_errors(self):
|
||
|
|
pyfunc = np_concatenate1
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
a = np.arange(15)
|
||
|
|
b = a.reshape((3, 5))
|
||
|
|
c = a.astype(np.dtype([('x', np.int8)]))
|
||
|
|
d = np.array(42)
|
||
|
|
|
||
|
|
# Different dimensionalities
|
||
|
|
with self.assertTypingError() as raises:
|
||
|
|
cfunc(a, b, b)
|
||
|
|
self.assertIn("all the input arrays must have same number of dimensions",
|
||
|
|
str(raises.exception))
|
||
|
|
|
||
|
|
# Incompatible dtypes
|
||
|
|
with self.assertTypingError() as raises:
|
||
|
|
cfunc(a, c, c)
|
||
|
|
self.assertIn("input arrays must have compatible dtypes",
|
||
|
|
str(raises.exception))
|
||
|
|
|
||
|
|
# 0-d arrays
|
||
|
|
with self.assertTypingError() as raises:
|
||
|
|
cfunc(d, d, d)
|
||
|
|
self.assertIn("zero-dimensional arrays cannot be concatenated",
|
||
|
|
str(raises.exception))
|
||
|
|
|
||
|
|
# non-tuple input
|
||
|
|
with self.assertTypingError() as raises:
|
||
|
|
cfunc(c, 1, c)
|
||
|
|
self.assertIn('expecting a non-empty tuple of arrays', str(raises.exception))
|
||
|
|
|
||
|
|
|
||
|
|
@unittest.skipUnless(hasattr(np, "stack"), "this Numpy doesn't have np.stack()")
|
||
|
|
class TestNpStack(MemoryLeakMixin, TestCase):
|
||
|
|
"""
|
||
|
|
Tests for np.stack().
|
||
|
|
"""
|
||
|
|
|
||
|
|
def _3d_arrays(self):
|
||
|
|
a = np.arange(24).reshape((4, 3, 2))
|
||
|
|
b = a + 10
|
||
|
|
c = (b + 10).copy(order='F')
|
||
|
|
d = (c + 10)[::-1]
|
||
|
|
e = (d + 10)[...,::-1]
|
||
|
|
return a, b, c, d, e
|
||
|
|
|
||
|
|
@contextlib.contextmanager
|
||
|
|
def assert_invalid_sizes(self):
|
||
|
|
with self.assertRaises(ValueError) as raises:
|
||
|
|
yield
|
||
|
|
self.assertIn("all input arrays must have the same shape",
|
||
|
|
str(raises.exception))
|
||
|
|
|
||
|
|
def check_stack(self, pyfunc, cfunc, args):
|
||
|
|
expected = pyfunc(*args)
|
||
|
|
got = cfunc(*args)
|
||
|
|
# Numba doesn't choose the same layout as Numpy.
|
||
|
|
# We would like to check the result is contiguous, but we can't
|
||
|
|
# rely on the "flags" attribute when there are 1-sized
|
||
|
|
# dimensions.
|
||
|
|
self.assertEqual(got.shape, expected.shape)
|
||
|
|
self.assertPreciseEqual(got.flatten(), expected.flatten())
|
||
|
|
|
||
|
|
def check_3d(self, pyfunc, cfunc, generate_starargs):
|
||
|
|
def check(a, b, c, args):
|
||
|
|
self.check_stack(pyfunc, cfunc, (a, b, c) + args)
|
||
|
|
|
||
|
|
def check_all_axes(a, b, c):
|
||
|
|
for args in generate_starargs():
|
||
|
|
check(a, b, c, args)
|
||
|
|
|
||
|
|
a, b, c, d, e = self._3d_arrays()
|
||
|
|
|
||
|
|
# C, C, C
|
||
|
|
check_all_axes(a, b, b)
|
||
|
|
# C, C, F
|
||
|
|
check_all_axes(a, b, c)
|
||
|
|
# F, F, F
|
||
|
|
check_all_axes(a.T, b.T, a.T)
|
||
|
|
# F, F, C
|
||
|
|
check_all_axes(a.T, b.T, c.T)
|
||
|
|
# F, F, A
|
||
|
|
check_all_axes(a.T, b.T, d.T)
|
||
|
|
# A, A, A
|
||
|
|
check_all_axes(d.T, e.T, d.T)
|
||
|
|
|
||
|
|
# Different but compatible dtypes
|
||
|
|
check_all_axes(a, b.astype(np.float64), b)
|
||
|
|
|
||
|
|
def check_runtime_errors(self, cfunc, generate_starargs):
|
||
|
|
# Exceptions leak references
|
||
|
|
self.assert_no_memory_leak()
|
||
|
|
self.disable_leak_check()
|
||
|
|
|
||
|
|
# Inputs have different shapes
|
||
|
|
a, b, c, d, e = self._3d_arrays()
|
||
|
|
with self.assert_invalid_sizes():
|
||
|
|
args = next(generate_starargs())
|
||
|
|
cfunc(a[:-1], b, c, *args)
|
||
|
|
|
||
|
|
def test_3d(self):
|
||
|
|
"""
|
||
|
|
stack(3d arrays, axis)
|
||
|
|
"""
|
||
|
|
pyfunc = np_stack2
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
def generate_starargs():
|
||
|
|
for axis in range(3):
|
||
|
|
yield (axis,)
|
||
|
|
yield (-3 + axis,)
|
||
|
|
|
||
|
|
self.check_3d(pyfunc, cfunc, generate_starargs)
|
||
|
|
self.check_runtime_errors(cfunc, generate_starargs)
|
||
|
|
|
||
|
|
def test_3d_no_axis(self):
|
||
|
|
"""
|
||
|
|
stack(3d arrays)
|
||
|
|
"""
|
||
|
|
pyfunc = np_stack1
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
def generate_starargs():
|
||
|
|
yield()
|
||
|
|
|
||
|
|
self.check_3d(pyfunc, cfunc, generate_starargs)
|
||
|
|
self.check_runtime_errors(cfunc, generate_starargs)
|
||
|
|
|
||
|
|
def test_0d(self):
|
||
|
|
"""
|
||
|
|
stack(0d arrays)
|
||
|
|
"""
|
||
|
|
pyfunc = np_stack1
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
a = np.array(42)
|
||
|
|
b = np.array(-5j)
|
||
|
|
c = np.array(True)
|
||
|
|
|
||
|
|
self.check_stack(pyfunc, cfunc, (a, b, c))
|
||
|
|
|
||
|
|
def check_xxstack(self, pyfunc, cfunc):
|
||
|
|
"""
|
||
|
|
3d and 0d tests for hstack(), vstack(), dstack().
|
||
|
|
"""
|
||
|
|
def generate_starargs():
|
||
|
|
yield()
|
||
|
|
|
||
|
|
self.check_3d(pyfunc, cfunc, generate_starargs)
|
||
|
|
# 0d
|
||
|
|
a = np.array(42)
|
||
|
|
b = np.array(-5j)
|
||
|
|
c = np.array(True)
|
||
|
|
self.check_stack(pyfunc, cfunc, (a, b, a))
|
||
|
|
|
||
|
|
def test_hstack(self):
|
||
|
|
pyfunc = np_hstack
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
self.check_xxstack(pyfunc, cfunc)
|
||
|
|
# 1d
|
||
|
|
a = np.arange(5)
|
||
|
|
b = np.arange(6) + 10
|
||
|
|
self.check_stack(pyfunc, cfunc, (a, b, b))
|
||
|
|
# 2d
|
||
|
|
a = np.arange(6).reshape((2, 3))
|
||
|
|
b = np.arange(8).reshape((2, 4)) + 100
|
||
|
|
self.check_stack(pyfunc, cfunc, (a, b, a))
|
||
|
|
|
||
|
|
def test_vstack(self):
|
||
|
|
# Since np.row_stack is an alias for np.vstack, it does not need a
|
||
|
|
# separate Numba implementation. For every test for np.vstack, the same
|
||
|
|
# test for np.row_stack has been added.
|
||
|
|
for pyfunc in (np_vstack, np_row_stack):
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
self.check_xxstack(pyfunc, cfunc)
|
||
|
|
# 1d
|
||
|
|
a = np.arange(5)
|
||
|
|
b = a + 10
|
||
|
|
self.check_stack(pyfunc, cfunc, (a, b, b))
|
||
|
|
# 2d
|
||
|
|
a = np.arange(6).reshape((3, 2))
|
||
|
|
b = np.arange(8).reshape((4, 2)) + 100
|
||
|
|
self.check_stack(pyfunc, cfunc, (a, b, b))
|
||
|
|
|
||
|
|
def test_dstack(self):
|
||
|
|
pyfunc = np_dstack
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
self.check_xxstack(pyfunc, cfunc)
|
||
|
|
# 1d
|
||
|
|
a = np.arange(5)
|
||
|
|
b = a + 10
|
||
|
|
self.check_stack(pyfunc, cfunc, (a, b, b))
|
||
|
|
# 2d
|
||
|
|
a = np.arange(12).reshape((3, 4))
|
||
|
|
b = a + 100
|
||
|
|
self.check_stack(pyfunc, cfunc, (a, b, b))
|
||
|
|
|
||
|
|
def test_column_stack(self):
|
||
|
|
pyfunc = np_column_stack
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
a = np.arange(4)
|
||
|
|
b = a + 10
|
||
|
|
c = np.arange(12).reshape((4, 3))
|
||
|
|
self.check_stack(pyfunc, cfunc, (a, b, c))
|
||
|
|
|
||
|
|
# Exceptions leak references
|
||
|
|
self.assert_no_memory_leak()
|
||
|
|
self.disable_leak_check()
|
||
|
|
|
||
|
|
# Invalid dims
|
||
|
|
a = np.array(42)
|
||
|
|
with self.assertTypingError():
|
||
|
|
cfunc((a, a, a))
|
||
|
|
a = a.reshape((1, 1, 1))
|
||
|
|
with self.assertTypingError():
|
||
|
|
cfunc((a, a, a))
|
||
|
|
|
||
|
|
def test_bad_arrays(self):
|
||
|
|
for pyfunc in (np_stack1, np_hstack, np_vstack, np_dstack, np_column_stack):
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
c = np.arange(12).reshape((4, 3))
|
||
|
|
|
||
|
|
# non-tuple input
|
||
|
|
with self.assertTypingError() as raises:
|
||
|
|
cfunc(c, 1, c)
|
||
|
|
self.assertIn('expecting a non-empty tuple of arrays', str(raises.exception))
|
||
|
|
|
||
|
|
|
||
|
|
def benchmark_refct_speed():
|
||
|
|
def pyfunc(x, y, t):
|
||
|
|
"""Swap array x and y for t number of times
|
||
|
|
"""
|
||
|
|
for i in range(t):
|
||
|
|
x, y = y, x
|
||
|
|
return x, y
|
||
|
|
|
||
|
|
cfunc = nrtjit(pyfunc)
|
||
|
|
|
||
|
|
x = np.random.random(100)
|
||
|
|
y = np.random.random(100)
|
||
|
|
t = 10000
|
||
|
|
|
||
|
|
def bench_pyfunc():
|
||
|
|
pyfunc(x, y, t)
|
||
|
|
|
||
|
|
def bench_cfunc():
|
||
|
|
cfunc(x, y, t)
|
||
|
|
|
||
|
|
python_time = utils.benchmark(bench_pyfunc)
|
||
|
|
numba_time = utils.benchmark(bench_cfunc)
|
||
|
|
print(python_time)
|
||
|
|
print(numba_time)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|