313 lines
9.5 KiB
Python
313 lines
9.5 KiB
Python
|
import itertools
|
||
|
|
import pickle
|
||
|
|
import textwrap
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from numba import njit, vectorize
|
||
|
|
from numba.tests.support import MemoryLeakMixin, TestCase
|
||
|
|
from numba.core.errors import TypingError
|
||
|
|
import unittest
|
||
|
|
from numba.np.ufunc import dufunc
|
||
|
|
|
||
|
|
|
||
|
|
def pyuadd(a0, a1):
|
||
|
|
return a0 + a1
|
||
|
|
|
||
|
|
|
||
|
|
def pysub(a0, a1):
|
||
|
|
return a0 - a1
|
||
|
|
|
||
|
|
|
||
|
|
def pymult(a0, a1):
|
||
|
|
return a0 * a1
|
||
|
|
|
||
|
|
|
||
|
|
def pydiv(a0, a1):
|
||
|
|
return a0 // a1
|
||
|
|
|
||
|
|
|
||
|
|
def pymin(a0, a1):
|
||
|
|
return a0 if a0 < a1 else a1
|
||
|
|
|
||
|
|
|
||
|
|
class TestDUFunc(MemoryLeakMixin, unittest.TestCase):
|
||
|
|
|
||
|
|
def nopython_dufunc(self, pyfunc):
|
||
|
|
return dufunc.DUFunc(pyfunc, targetoptions=dict(nopython=True))
|
||
|
|
|
||
|
|
def test_frozen(self):
|
||
|
|
duadd = self.nopython_dufunc(pyuadd)
|
||
|
|
self.assertFalse(duadd._frozen)
|
||
|
|
duadd._frozen = True
|
||
|
|
self.assertTrue(duadd._frozen)
|
||
|
|
with self.assertRaises(ValueError):
|
||
|
|
duadd._frozen = False
|
||
|
|
with self.assertRaises(TypeError):
|
||
|
|
duadd(np.linspace(0,1,10), np.linspace(1,2,10))
|
||
|
|
|
||
|
|
def test_scalar(self):
|
||
|
|
duadd = self.nopython_dufunc(pyuadd)
|
||
|
|
self.assertEqual(pyuadd(1,2), duadd(1,2))
|
||
|
|
|
||
|
|
def test_npm_call(self):
|
||
|
|
duadd = self.nopython_dufunc(pyuadd)
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def npmadd(a0, a1, o0):
|
||
|
|
duadd(a0, a1, o0)
|
||
|
|
X = np.linspace(0,1.9,20)
|
||
|
|
X0 = X[:10]
|
||
|
|
X1 = X[10:]
|
||
|
|
out0 = np.zeros(10)
|
||
|
|
npmadd(X0, X1, out0)
|
||
|
|
np.testing.assert_array_equal(X0 + X1, out0)
|
||
|
|
Y0 = X0.reshape((2,5))
|
||
|
|
Y1 = X1.reshape((2,5))
|
||
|
|
out1 = np.zeros((2,5))
|
||
|
|
npmadd(Y0, Y1, out1)
|
||
|
|
np.testing.assert_array_equal(Y0 + Y1, out1)
|
||
|
|
Y2 = X1[:5]
|
||
|
|
out2 = np.zeros((2,5))
|
||
|
|
npmadd(Y0, Y2, out2)
|
||
|
|
np.testing.assert_array_equal(Y0 + Y2, out2)
|
||
|
|
|
||
|
|
def test_npm_call_implicit_output(self):
|
||
|
|
duadd = self.nopython_dufunc(pyuadd)
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def npmadd(a0, a1):
|
||
|
|
return duadd(a0, a1)
|
||
|
|
X = np.linspace(0,1.9,20)
|
||
|
|
X0 = X[:10]
|
||
|
|
X1 = X[10:]
|
||
|
|
out0 = npmadd(X0, X1)
|
||
|
|
np.testing.assert_array_equal(X0 + X1, out0)
|
||
|
|
Y0 = X0.reshape((2,5))
|
||
|
|
Y1 = X1.reshape((2,5))
|
||
|
|
out1 = npmadd(Y0, Y1)
|
||
|
|
np.testing.assert_array_equal(Y0 + Y1, out1)
|
||
|
|
Y2 = X1[:5]
|
||
|
|
out2 = npmadd(Y0, Y2)
|
||
|
|
np.testing.assert_array_equal(Y0 + Y2, out2)
|
||
|
|
out3 = npmadd(1.,2.)
|
||
|
|
self.assertEqual(out3, 3.)
|
||
|
|
|
||
|
|
def test_ufunc_props(self):
|
||
|
|
duadd = self.nopython_dufunc(pyuadd)
|
||
|
|
self.assertEqual(duadd.nin, 2)
|
||
|
|
self.assertEqual(duadd.nout, 1)
|
||
|
|
self.assertEqual(duadd.nargs, duadd.nin + duadd.nout)
|
||
|
|
self.assertEqual(duadd.ntypes, 0)
|
||
|
|
self.assertEqual(duadd.types, [])
|
||
|
|
self.assertEqual(duadd.identity, None)
|
||
|
|
duadd(1, 2)
|
||
|
|
self.assertEqual(duadd.ntypes, 1)
|
||
|
|
self.assertEqual(duadd.ntypes, len(duadd.types))
|
||
|
|
self.assertIsNone(duadd.signature)
|
||
|
|
|
||
|
|
def test_ufunc_props_jit(self):
|
||
|
|
duadd = self.nopython_dufunc(pyuadd)
|
||
|
|
duadd(1, 2) # initialize types attribute
|
||
|
|
|
||
|
|
attributes = {'nin': duadd.nin,
|
||
|
|
'nout': duadd.nout,
|
||
|
|
'nargs': duadd.nargs,
|
||
|
|
#'ntypes': duadd.ntypes,
|
||
|
|
#'types': duadd.types,
|
||
|
|
'identity': duadd.identity,
|
||
|
|
'signature': duadd.signature}
|
||
|
|
|
||
|
|
def get_attr_fn(attr):
|
||
|
|
fn = f'''
|
||
|
|
def impl():
|
||
|
|
return duadd.{attr}
|
||
|
|
'''
|
||
|
|
l = {}
|
||
|
|
exec(textwrap.dedent(fn), {'duadd': duadd}, l)
|
||
|
|
return l['impl']
|
||
|
|
|
||
|
|
for attr, val in attributes.items():
|
||
|
|
cfunc = njit(get_attr_fn(attr))
|
||
|
|
self.assertEqual(val, cfunc(),
|
||
|
|
f'Attribute differs from original: {attr}')
|
||
|
|
|
||
|
|
# We don't expose [n]types attributes as they are dynamic attributes
|
||
|
|
# and can change as the user calls the ufunc
|
||
|
|
# cfunc = njit(get_attr_fn('ntypes'))
|
||
|
|
# self.assertEqual(cfunc(), 1)
|
||
|
|
# duadd(1.1, 2.2)
|
||
|
|
# self.assertEqual(cfunc(), 2)
|
||
|
|
|
||
|
|
|
||
|
|
class TestDUFuncMethods(TestCase):
|
||
|
|
def _check_reduce(self, ufunc, dtype=None, initial=None):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(a, axis, dtype, initial):
|
||
|
|
return ufunc.reduce(a,
|
||
|
|
axis=axis,
|
||
|
|
dtype=dtype,
|
||
|
|
initial=initial)
|
||
|
|
|
||
|
|
inputs = [
|
||
|
|
np.arange(5),
|
||
|
|
np.arange(4).reshape(2, 2),
|
||
|
|
np.arange(40).reshape(5, 4, 2),
|
||
|
|
]
|
||
|
|
for array in inputs:
|
||
|
|
for axis in range(array.ndim):
|
||
|
|
expected = foo.py_func(array, axis, dtype, initial)
|
||
|
|
got = foo(array, axis, dtype, initial)
|
||
|
|
self.assertPreciseEqual(expected, got)
|
||
|
|
|
||
|
|
def _check_reduce_axis(self, ufunc, dtype, initial=None):
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(a, axis):
|
||
|
|
return ufunc.reduce(a, axis=axis, initial=initial)
|
||
|
|
|
||
|
|
def _check(*args):
|
||
|
|
try:
|
||
|
|
expected = foo.py_func(array, axis)
|
||
|
|
except ValueError as e:
|
||
|
|
self.assertEqual(e.args[0], exc_msg)
|
||
|
|
with self.assertRaisesRegex(TypingError, exc_msg):
|
||
|
|
got = foo(array, axis)
|
||
|
|
else:
|
||
|
|
got = foo(array, axis)
|
||
|
|
self.assertPreciseEqual(expected, got)
|
||
|
|
|
||
|
|
exc_msg = (f"reduction operation '{ufunc.__name__}' is not "
|
||
|
|
"reorderable, so at most one axis may be specified")
|
||
|
|
inputs = [
|
||
|
|
np.arange(40, dtype=dtype).reshape(5, 4, 2),
|
||
|
|
np.arange(10, dtype=dtype),
|
||
|
|
]
|
||
|
|
for array in inputs:
|
||
|
|
for i in range(1, array.ndim + 1):
|
||
|
|
for axis in itertools.combinations(range(array.ndim), r=i):
|
||
|
|
_check(array, axis)
|
||
|
|
|
||
|
|
# corner cases: Reduce over axis=() and axis=None
|
||
|
|
for axis in ((), None):
|
||
|
|
_check(array, axis)
|
||
|
|
|
||
|
|
def test_add_reduce(self):
|
||
|
|
duadd = vectorize('int64(int64, int64)', identity=0)(pyuadd)
|
||
|
|
self._check_reduce(duadd)
|
||
|
|
self._check_reduce_axis(duadd, dtype=np.int64)
|
||
|
|
|
||
|
|
def test_mul_reduce(self):
|
||
|
|
dumul = vectorize('int64(int64, int64)', identity=1)(pymult)
|
||
|
|
self._check_reduce(dumul)
|
||
|
|
|
||
|
|
def test_non_associative_reduce(self):
|
||
|
|
dusub = vectorize('int64(int64, int64)')(pysub)
|
||
|
|
dudiv = vectorize('int64(int64, int64)')(pydiv)
|
||
|
|
self._check_reduce(dusub)
|
||
|
|
self._check_reduce_axis(dusub, dtype=np.int64)
|
||
|
|
self._check_reduce(dudiv)
|
||
|
|
self._check_reduce_axis(dudiv, dtype=np.int64)
|
||
|
|
|
||
|
|
def test_reduce_dtype(self):
|
||
|
|
duadd = vectorize('float64(float64, int64)', identity=0)(pyuadd)
|
||
|
|
self._check_reduce(duadd, dtype=np.float64)
|
||
|
|
|
||
|
|
def test_min_reduce(self):
|
||
|
|
dumin = vectorize('int64(int64, int64)')(pymin)
|
||
|
|
self._check_reduce(dumin, initial=10)
|
||
|
|
self._check_reduce_axis(dumin, dtype=np.int64)
|
||
|
|
|
||
|
|
def test_add_reduce_initial(self):
|
||
|
|
# Initial should be used as a start
|
||
|
|
duadd = vectorize('int64(int64, int64)', identity=0)(pyuadd)
|
||
|
|
self._check_reduce(duadd, dtype=np.int64, initial=100)
|
||
|
|
|
||
|
|
def test_add_reduce_no_initial_or_identity(self):
|
||
|
|
# don't provide an initial or identity value
|
||
|
|
duadd = vectorize('int64(int64, int64)')(pyuadd)
|
||
|
|
self._check_reduce(duadd, dtype=np.int64)
|
||
|
|
|
||
|
|
def test_invalid_input(self):
|
||
|
|
duadd = vectorize('float64(float64, int64)', identity=0)(pyuadd)
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(a):
|
||
|
|
return duadd.reduce(a)
|
||
|
|
|
||
|
|
exc_msg = 'The first argument "array" must be array-like'
|
||
|
|
with self.assertRaisesRegex(TypingError, exc_msg):
|
||
|
|
foo('a')
|
||
|
|
|
||
|
|
def test_dufunc_negative_axis(self):
|
||
|
|
duadd = vectorize('int64(int64, int64)', identity=0)(pyuadd)
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(a, axis):
|
||
|
|
return duadd.reduce(a, axis=axis)
|
||
|
|
|
||
|
|
a = np.arange(40).reshape(5, 4, 2)
|
||
|
|
cases = (0, -1, (0, -1), (-1, -2), (1, -1), -3)
|
||
|
|
for axis in cases:
|
||
|
|
expected = duadd.reduce(a, axis)
|
||
|
|
got = foo(a, axis)
|
||
|
|
self.assertPreciseEqual(expected, got)
|
||
|
|
|
||
|
|
def test_dufunc_invalid_axis(self):
|
||
|
|
duadd = vectorize('int64(int64, int64)', identity=0)(pyuadd)
|
||
|
|
|
||
|
|
@njit
|
||
|
|
def foo(a, axis):
|
||
|
|
return duadd.reduce(a, axis=axis)
|
||
|
|
|
||
|
|
a = np.arange(40).reshape(5, 4, 2)
|
||
|
|
cases = ((0, 0), (0, 1, 0), (0, -3), (-1, -1), (-1, 2))
|
||
|
|
for axis in cases:
|
||
|
|
msg = "duplicate value in 'axis'"
|
||
|
|
with self.assertRaisesRegex(ValueError, msg):
|
||
|
|
foo(a, axis)
|
||
|
|
|
||
|
|
cases = (-4, 3, (0, -4),)
|
||
|
|
for axis in cases:
|
||
|
|
with self.assertRaisesRegex(ValueError, "Invalid axis"):
|
||
|
|
foo(a, axis)
|
||
|
|
|
||
|
|
|
||
|
|
class TestDUFuncPickling(MemoryLeakMixin, unittest.TestCase):
|
||
|
|
def check(self, ident, result_type):
|
||
|
|
buf = pickle.dumps(ident)
|
||
|
|
rebuilt = pickle.loads(buf)
|
||
|
|
|
||
|
|
# Check reconstructed dufunc
|
||
|
|
r = rebuilt(123)
|
||
|
|
self.assertEqual(123, r)
|
||
|
|
self.assertIsInstance(r, result_type)
|
||
|
|
|
||
|
|
# Try to use reconstructed dufunc in @jit
|
||
|
|
@njit
|
||
|
|
def foo(x):
|
||
|
|
return rebuilt(x)
|
||
|
|
|
||
|
|
r = foo(321)
|
||
|
|
self.assertEqual(321, r)
|
||
|
|
self.assertIsInstance(r, result_type)
|
||
|
|
|
||
|
|
def test_unrestricted(self):
|
||
|
|
@vectorize
|
||
|
|
def ident(x1):
|
||
|
|
return x1
|
||
|
|
|
||
|
|
self.check(ident, result_type=(int, np.integer))
|
||
|
|
|
||
|
|
def test_restricted(self):
|
||
|
|
@vectorize(["float64(float64)"])
|
||
|
|
def ident(x1):
|
||
|
|
return x1
|
||
|
|
|
||
|
|
self.check(ident, result_type=float)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|