def _ufunc_lookup(cls, name): """ Returns an arithmetic ufunc using lookup tables. These ufuncs are compiled for each Galois field since the lookup tables are compiled into the ufuncs as constants. """ key = (name, cls.characteristic, cls.degree, cls._irreducible_poly_int) if key not in cls._UFUNC_CACHE_LOOKUP: # These variables must be locals for Numba to compile them as literals EXP = cls._EXP LOG = cls._LOG ZECH_LOG = cls._ZECH_LOG ZECH_E = cls._ZECH_E function = getattr(cls, f"_{name}_lookup") if cls._UFUNC_TYPE[name] == "unary": cls._UFUNC_CACHE_LOOKUP[key] = numba.vectorize( ["int64(int64)"], nopython=True)( lambda a: function(a, EXP, LOG, ZECH_LOG, ZECH_E)) else: cls._UFUNC_CACHE_LOOKUP[key] = numba.vectorize( ["int64(int64, int64)"], nopython=True)(lambda a, b: function( a, b, EXP, LOG, ZECH_LOG, ZECH_E)) return cls._UFUNC_CACHE_LOOKUP[key]
def test_vectorize_no_args(self): a = numpy.linspace(0, 1, 10) b = numpy.linspace(1, 2, 10) ufunc = vectorize(add) self.assertTrue(numpy.all(ufunc(a, b) == (a + b))) ufunc2 = vectorize(add) c = numpy.empty(10) ufunc2(a, b, c) self.assertTrue(numpy.all(c == (a + b)))
def test_vectorize_no_args(self): a = numpy.linspace(0,1,10) b = numpy.linspace(1,2,10) ufunc = vectorize(add) self.assertTrue(numpy.all(ufunc(a,b) == (a + b))) ufunc2 = vectorize(add) c = numpy.empty(10) ufunc2(a, b, c) self.assertTrue(numpy.all(c == (a + b)))
def test_vectorize_no_args(self): a = np.linspace(0,1,10) b = np.linspace(1,2,10) ufunc = vectorize(add) self.assertPreciseEqual(ufunc(a,b), a + b) ufunc2 = vectorize(add) c = np.empty(10) ufunc2(a, b, c) self.assertPreciseEqual(c, a + b)
def test_is_jitted(self): def foo(x): pass self.assertFalse(is_jitted(foo)) self.assertTrue(is_jitted(njit(foo))) self.assertFalse(is_jitted(vectorize(foo))) self.assertFalse(is_jitted(vectorize(parallel=True)(foo))) self.assertFalse(is_jitted(guvectorize("void(float64[:])", "(m)")(foo)))
def test_cpu_vs_parallel(self): @jit def add(x, y): return x + y custom_vectorize = vectorize([], identity=None, target='cpu') custom_vectorize(add) custom_vectorize_2 = vectorize([], identity=None, target='parallel') custom_vectorize_2(add)
def _compile_lookup_ufuncs(cls, target): # Export lookup tables to global variables so JIT compiling can cache the tables in the binaries global CHARACTERISTIC, ORDER, EXP, LOG, ZECH_LOG, ZECH_E, ADD_JIT, MULTIPLY_JIT # pylint: disable=global-statement CHARACTERISTIC = cls.characteristic ORDER = cls.order EXP = cls._EXP LOG = cls._LOG ZECH_LOG = cls._ZECH_LOG if cls.characteristic == 2: ZECH_E = 0 else: ZECH_E = (cls.order - 1) // 2 kwargs = {"nopython": True, "target": target} if target == "cuda": kwargs.pop("nopython") # JIT-compile add and multiply routines for reference in other routines ADD_JIT = numba.jit("int64(int64, int64)", nopython=True)(_add_lookup) MULTIPLY_JIT = numba.jit("int64(int64, int64)", nopython=True)(_multiply_lookup) # Create numba JIT-compiled ufuncs using the *current* EXP, LOG, and MUL_INV lookup tables cls._numba_ufunc_add = numba.vectorize(["int64(int64, int64)"], **kwargs)(_add_lookup) cls._numba_ufunc_subtract = numba.vectorize(["int64(int64, int64)"], **kwargs)(_subtract_lookup) cls._numba_ufunc_multiply = numba.vectorize(["int64(int64, int64)"], **kwargs)(_multiply_lookup) cls._numba_ufunc_divide = numba.vectorize(["int64(int64, int64)"], **kwargs)(_divide_lookup) cls._numba_ufunc_negative = numba.vectorize(["int64(int64)"], **kwargs)(_additive_inverse_lookup) cls._numba_ufunc_multiple_add = numba.vectorize(["int64(int64, int64)"], **kwargs)(_multiple_add_lookup) cls._numba_ufunc_power = numba.vectorize(["int64(int64, int64)"], **kwargs)(_power_lookup) cls._numba_ufunc_log = numba.vectorize(["int64(int64)"], **kwargs)(_log_lookup) cls._numba_ufunc_poly_eval = numba.guvectorize([(numba.int64[:], numba.int64[:], numba.int64[:])], "(n),(m)->(m)", **kwargs)(_poly_eval_lookup)
def __init__(self, x_max, dx, dt, potential=None, stationary=False): """ Creates a solver for the interval [-x_max, x_max] with a spatial step `dx` and a time step `dt`. :param x_max: :param dx: :param dt: :param potential: :param stationary: if True, the potential is considered constant in time """ self.dx = dx self.dt = dt n1 = int(x_max / dx) if n1 * dx < x_max - 0.01 * dx: n1 += 1 self.n_points = 2 * n1 + 1 self.x = _np.linspace(-n1 * dx, n1 * dx, self.n_points) self.stationary = stationary if potential is None: self.potential = _numba.vectorize(nopython=True)(lambda x: 0.0j) self.stationary = True elif isinstance(potential, _potential.Potential): if not isinstance(potential, _potential.Potential1D): raise ValueError("Only 1D potentials are allowed") self.potential = potential.get_potential() self.stationary = potential.is_stationary() self.delta_depth = potential.get_delta_depth() else: self.potential = potential self._psi = _np.zeros(self.x.shape, dtype=_np.complex)
def numba_funcify_Elemwise(op, node, use_signature=False, identity=None, **kwargs): scalar_op_fn = numba_funcify(op.scalar_op, node, **kwargs) input_names = ", ".join([v.auto_name for v in node.inputs]) if use_signature: signature = [create_numba_signature(node, force_scalar=True)] else: signature = [] numba_vectorize = numba.vectorize(signature, identity=identity) global_env = { "scalar_op": scalar_op_fn, "numba_vectorize": numba_vectorize } elemwise_fn_name = f"elemwise_{get_name_for_object(scalar_op_fn)}" elemwise_src = f""" @numba_vectorize def {elemwise_fn_name}({input_names}): return scalar_op({input_names}) """ elemwise_fn = compile_function_src(elemwise_src, elemwise_fn_name, global_env) return elemwise_fn
def __init__(self, expression, arguments, argument_dtypes, return_dtype, verbose=False): self.expression = expression self.arguments = arguments self.argument_dtypes = argument_dtypes self.return_dtype = return_dtype self.verbose = verbose import numba argument_dtypes_numba = [ getattr(numba, argument_dtype.name) for argument_dtype in argument_dtypes ] argstring = ", ".join(arguments) code = ''' from numpy import * def f({0}): return {1}'''.format(argstring, expression) if verbose: print('Generated code:\n' + code) scope = {} exec(code, scope) f = scope['f'] return_dtype_numba = getattr(numba, return_dtype.name) vectorizer = numba.vectorize( [return_dtype_numba(*argument_dtypes_numba)]) self.f = vectorizer(f)
def main(): N = 1e+8 // 2 print('Data size', N) targets = ['cpu', 'parallel'] # run just one target if is specified in the argument for t in targets: if t in sys.argv[1:]: targets = [t] break for target in targets: print('== Target', target) vect_discriminant = vectorize(['f4(f4, f4, f4)', 'f8(f8, f8, f8)'], target=target)(discriminant) A, B, C = generate_input(N, dtype=np.float32) D = np.empty(A.shape, dtype=A.dtype) ts = time() D = vect_discriminant(A, B, C) te = time() total_time = (te - ts) print('Execution time %.4f' % total_time) print('Throughput %.4f' % (N / total_time)) if '-verify' in sys.argv[1:]: check_answer(D, A, B, C)
def main(): N = 1e8 // 2 print("Data size", N) targets = ["cpu", "parallel"] # run just one target if is specified in the argument for t in targets: if t in sys.argv[1:]: targets = [t] break for target in targets: print("== Target", target) vect_discriminant = vectorize(["f4(f4, f4, f4)", "f8(f8, f8, f8)"], target=target)(discriminant) A, B, C = generate_input(N, dtype=np.float32) D = np.empty(A.shape, dtype=A.dtype) ts = time() D = vect_discriminant(A, B, C) te = time() total_time = te - ts print("Execution time %.4f" % total_time) print("Throughput %.4f" % (N / total_time)) if "-verify" in sys.argv[1:]: check_answer(D, A, B, C)
def main(): targets = ['cpu', 'parallel'] # run just one target if is specified in the argument for t in targets: if t in sys.argv[1:]: targets = [t] break for target in targets: print('== Target', target) vect_sum = vectorize(['f4(f4, f4)', 'f8(f8, f8)'], target=target)(sum) A = np.random.random(N).astype(np.float32) B = np.random.random(N).astype(np.float32) assert A.shape == B.shape assert A.dtype == B.dtype assert len(A.shape) == 1 D = np.empty(A.shape, dtype=A.dtype) print('Data size', N) ts = time() D = vect_sum(A, B) te = time() total_time = (te - ts) print('Execution time %.4f' % total_time) print('Throughput %.4f' % (N / total_time)) if '-verify' in sys.argv[1:]: check_answer(D, A, B, C)
def greyscale_image(arr): vectorized_greyscale_img = nb.vectorize(greyscale_pixel) split_arr = np.split(arr, 3) res = vectorized_greyscale_img(split_arr[0], split_arr[1], split_arr[2]) return res
def __apply_MapOp(mesh, particles_pos, shape_function, mapOp): half_supp = 0.5 * shape_function.support num_coeffs_per_axis = int(np.ceil(shape_function.support)) dim = len(np.shape(mesh)) # ndarray which holds the non-zero particle assignment values of one particle coeffs = np.empty([num_coeffs_per_axis] * dim) if vectorize: # vectorize shape function to a compiled ufunc v_shape_function = vectorize(['f4(f4)', 'f8(f8)'])(shape_function.__call__) else: # numpy fallback v_shape_function = np.vectorize(shape_function.__call__) for idx, pos in enumerate(particles_pos): # lowest mesh indices that contributes to the mapping begin = np.ceil(pos - np.ones(dim) * half_supp).astype(int) # rearrange mapping area if it overlaps the border begin = np.maximum(begin, np.zeros(dim, dtype=int)) begin = np.minimum(begin, np.shape(mesh) - np.ones(dim) * half_supp) # compute coefficients (particle assignment values) for coeff_idx in np.ndindex(np.shape(coeffs)): rel_vec = begin + coeff_idx - pos coeffs[coeff_idx] = np.prod(v_shape_function(rel_vec)) # do the actual mapping window = [slice(begin[i], begin[i] + num_coeffs_per_axis) for i in range(dim)] mapOp(mesh[window], coeffs, idx)
def __apply_MapOp(mesh, particles_pos, shape_function, mapOp): half_supp = 0.5 * shape_function.support num_coeffs_per_axis = int(np.ceil(shape_function.support)) dim = len(np.shape(mesh)) # ndarray which holds the non-zero particle assignment values of one particle coeffs = np.empty([num_coeffs_per_axis] * dim) if vectorize: # vectorize shape function to a compiled ufunc v_shape_function = vectorize(['f4(f4)', 'f8(f8)'])(shape_function.__call__) else: # numpy fallback v_shape_function = np.vectorize(shape_function.__call__) for idx, pos in enumerate(particles_pos): # lowest mesh indices that contributes to the mapping begin = np.ceil(pos - np.ones(dim) * half_supp).astype(int) # rearrange mapping area if it overlaps the border begin = np.maximum(begin, np.zeros(dim, dtype=int)) begin = np.minimum(begin, np.shape(mesh) - np.ones(dim) * half_supp) # compute coefficients (particle assignment values) for coeff_idx in np.ndindex(np.shape(coeffs)): rel_vec = begin + coeff_idx - pos coeffs[coeff_idx] = np.prod(v_shape_function(rel_vec)) # do the actual mapping window = [ slice(begin[i], begin[i] + num_coeffs_per_axis) for i in range(dim) ] mapOp(mesh[window], coeffs, idx)
def operator(f): """Define and register a new composite operator""" f2 = nb.vectorize(f) f2._compile_for_argtys((nb.types.uint32, nb.types.uint32)) f2._frozen = True composite_op_lookup[f.__name__] = f2 return f2
def test_vectorize(filter_str, shape, dtypes, input_type): if _helper.platform_not_supported(filter_str): pytest.skip() if _helper.skip_test(filter_str): pytest.skip() def vector_add(a, b): return a + b dtype, sig_dtype = dtypes sig = [sig_dtype(sig_dtype, sig_dtype)] size, shape = shape if input_type == "array": A = np.arange(size, dtype=dtype).reshape(shape) B = np.arange(size, dtype=dtype).reshape(shape) elif input_type == "scalar": A = dtype(1.2) B = dtype(2.3) with dppy.offload_to_sycl_device(filter_str): f = vectorize(sig, target="dppy")(vector_add) expected = f(A, B) actual = vector_add(A, B) max_abs_err = np.sum(expected) - np.sum(actual) assert max_abs_err < 1e-5
def _run_and_compare(cls, func, sig, A, *args, **kwargs): if cls.wrapper is not None: func = cls.wrapper(func) numba_func = vectorize(sig, target=cls.target)(func) numpy_func = np.vectorize(func) result = numba_func(A, *args) gold = numpy_func(A, *args) np.testing.assert_allclose(result, gold, **kwargs)
def _test_target_nopython(self, target, warnings, with_sig=True): a = np.array([2.0], dtype=np.float32) b = np.array([3.0], dtype=np.float32) sig = [float32(float32, float32)] args = with_sig and [sig] or [] with self.check_warnings(warnings): f = vectorize(*args, target=target, nopython=True)(vector_add) f(a, b)
def _compile(self, expr, args, argtypes, function_name): function_str = '''def %s(%s): return %s ''' % (function_name, ','.join(args), expr) scope = {_implicit_math_module: math} exec(function_str, scope) vectorizer = vectorize([argtypes], target=self.target) return vectorizer(scope[function_name])
def _gmf_function(self, ftype='numba_vectorize'): """ get vectorized function for gmf Parameters ---------- name: str gmf name ftype: str return function type. Allowed values are: - 'numba_vectorize': vectorized with `numba.vectorize` (default) - 'numba_guvectorize': vectorized with `numba.guvectorize` - 'vectorize': vectorized with `numpy.vectorize`, without signature - 'guvectorize': vectorized with `numpy.vectorize`, with signature - 'numba_njit': compiled with `numba.njit` (only scalar input will be allowed) - None: pure python function (only scalar input will be allowed) Returns ------- function `sigma0_linear = function(inc, wspd, [phi])` """ gmf_function = None if ftype is None: gmf_function = self._gmf_pyfunc_scalar elif ftype == 'numba_njit': gmf_function = njit([float64(float64, float64, float64)], nogil=True, inline='never')( self._gmf_pyfunc_scalar) elif ftype == 'numba_vectorize': gmf_function = vectorize( [ float64(float64, float64, float64), float32(float32, float32, float64) ], target='parallel', nopython=True)(self._gmf_pyfunc_scalar) elif ftype == 'numba_guvectorize': func_njit = self._gmf_function(ftype='numba_njit') @guvectorize( [ (float64[:], float64[:], float64[:], float64[:, :, :]), (float32[:], float32[:], float32[:], float32[:, :, :]) ], '(n),(m),(p)->(n,m,p)', target='cpu') def func(inc, wspd, phi, sigma0_out): for i_phi, one_phi in enumerate(phi): for i_wspd, one_wspd in enumerate(wspd): for i_inc, one_inc in enumerate(inc): sigma0_out[i_inc, i_wspd, i_phi] = func_njit(one_inc, one_wspd, one_phi) gmf_function = func else: raise TypeError('ftype "%s" not known' % ftype) return gmf_function
def _test_template_3(self, target): numba_scaled_sinc = vectorize(['float64(float64, uint32)'], target=target)(scaled_sinc) numpy_scaled_sinc = np.vectorize(scaled_sinc) A = np.arange(100, dtype=np.float64) scale = np.uint32(3) result = numba_scaled_sinc(A, scale) gold = numpy_scaled_sinc(A, scale) np.testing.assert_allclose(result, gold, atol=1e-8)
def check_ufunc_raise(self, **vectorize_args): f = vectorize(['float64(float64)'], **vectorize_args)(sqrt) arr = np.array([1, 4, -2, 9, -1, 16], dtype=np.float64) out = np.zeros_like(arr) with self.assertRaises(ValueError) as cm: f(arr, out) self.assertIn('Value must be positive', str(cm.exception)) # All values were computed except for the ones giving an error self.assertEqual(list(out), [1, 2, 0, 3, 0, 4])
def compile_func_combinations(func_combinations): d = {} for fc in func_combinations: s, t = sympy.symbols("s t") expr_temp = func_combinations[fc](s, t) fn = lambdify((s, t), expr_temp) vect = nb.vectorize(["float32(float32, float32)"], nopython=True) d[fc] = vect(fn) return d
def _test_target_unrecognized_arg(self, target, with_sig=True): a = np.array([2.0], dtype=np.float32) b = np.array([3.0], dtype=np.float32) sig = [float32(float32, float32)] args = with_sig and [sig] or [] with self.assertRaises(KeyError) as raises: f = vectorize(*args, target=target, nonexistent=2)(vector_add) f(a, b) self.assertIn("Unrecognized options", str(raises.exception))
def make_wrapper(func): vecd_f = vectorize(dtype_args)(func) @wraps(func) def wrapper(*args, **kwargs): arrays = [arg.values for arg in args] ans = vecd_f(*arrays) return ans return wrapper
def test_vectorize(self): def foo(x): return x + math.sin(x) fastfoo = vectorize(fastmath=True)(foo) slowfoo = vectorize(foo) x = np.random.random(8).astype(np.float32) # capture the optimized llvm to check for fast flag with override_config('DUMP_OPTIMIZED', True): with captured_stdout() as slow_cap: expect = slowfoo(x) slowllvm = slow_cap.getvalue() with captured_stdout() as fast_cap: got = fastfoo(x) fastllvm = fast_cap.getvalue() np.testing.assert_almost_equal(expect, got) self.assertIn('fadd fast', fastllvm) self.assertIn('call fast', fastllvm) self.assertNotIn('fadd fast', slowllvm) self.assertNotIn('call fast', slowllvm)
def test_vectorize_output_kwarg(self): """ Passing the output array as a keyword argument (issue #1867). """ def check(ufunc): a = np.arange(10, 16, dtype='int32') out = np.zeros_like(a) got = ufunc(a, a, out=out) self.assertIs(got, out) self.assertPreciseEqual(out, a + a) with self.assertRaises(TypeError): ufunc(a, a, zzz=out) # With explicit sigs ufunc = vectorize(['int32(int32, int32)'], nopython=True)(add) check(ufunc) # With implicit sig ufunc = vectorize(nopython=True)(add) check(ufunc) # compiling check(ufunc) # after compiling
def check_divmod_float(self, pyfunc, values, messages): """ Test 1 // 0 and 0 // 0. """ f = vectorize(nopython=True)(pyfunc) a = np.array([5., 6., 0., 9.]) b = np.array([1., 0., 0., 4.]) expected = np.array(values) with self.check_warnings(messages): res = f(a, b) self.assertPreciseEqual(res, expected)
def check_divmod_int(self, pyfunc, values): """ Test 1 % 0 and 0 % 0. """ f = vectorize(nopython=True)(pyfunc) a = np.array([5, 6, 0, 9]) b = np.array([1, 0, 0, 4]) expected = np.array(values) # No warnings raised because LLVM makes it difficult with self.check_warnings([]): res = f(a, b) self.assertPreciseEqual(res, expected)
def test_power_float(self): """ Test 0 ** -1 and 2 ** <big number>. """ f = vectorize(nopython=True)(power) a = np.array([5., 0., 2., 8.]) b = np.array([1., -1., 1e20, 4.]) expected = np.array([5., float('inf'), float('inf'), 4096.]) with self.check_warnings(["divide by zero encountered", "overflow encountered"]): res = f(a, b) self.assertPreciseEqual(res, expected)
def check_truediv_real(self, dtype): """ Test 1 / 0 and 0 / 0. """ f = vectorize(nopython=True)(truediv) a = np.array([5., 6., 0., 8.], dtype=dtype) b = np.array([1., 0., 0., 4.], dtype=dtype) expected = np.array([5., float('inf'), float('nan'), 2.]) with self.check_warnings(["divide by zero encountered", "invalid value encountered"]): res = f(a, b) self.assertPreciseEqual(res, expected)
def check_truediv_real(self, dtype): """ Test 1 / 0 and 0 / 0. """ f = vectorize(nopython=True)(truediv) a = np.array([5., 6., 0., 8.], dtype=dtype) b = np.array([1., 0., 0., 4.], dtype=dtype) expected = np.array([5., float('inf'), float('nan'), 2.]) with self.check_warnings( ["divide by zero encountered", "invalid value encountered"]): res = f(a, b) self.assertPreciseEqual(res, expected)
def test_power_float(self): """ Test 0 ** -1 and 2 ** <big number>. """ f = vectorize(nopython=True)(power) a = np.array([5., 0., 2., 8.]) b = np.array([1., -1., 1e20, 4.]) expected = np.array([5., float('inf'), float('inf'), 4096.]) with self.check_warnings( ["divide by zero encountered", "overflow encountered"]): res = f(a, b) self.assertPreciseEqual(res, expected)
def _get_numba_ufunc(expr): """Construct a numba ufunc from a blaze expression Parameters ---------- expr : blaze.expr.Expr Returns ------- f : function A numba vectorized function Examples -------- >>> from blaze import symbol >>> import numpy as np >>> s = symbol('s', 'float64') >>> t = symbol('t', 'float64') >>> x = np.array([1.0, 2.0, 3.0]) >>> y = np.array([2.0, 3.0, 4.0]) >>> f = get_numba_ufunc(s + t) >>> f(x, y) array([ 3., 5., 7.]) See Also -------- get_numba_type compute_signature """ if isinstance(expr, Broadcast): leaves = expr._scalars expr = expr._scalar_expr else: leaves = expr._leaves() s, scope = funcstr(leaves, expr) scope = dict((k, numba.jit(nopython=True)(v) if callable(v) else v) for k, v in scope.items()) # get the func func = eval(s, scope) # get the signature sig = compute_signature(expr) # vectorize is currently not thread safe. So lock the thread. # TODO FIXME remove this when numba has made vectorize thread safe. with lock: ufunc = numba.vectorize([sig], nopython=True)(func) return ufunc
def lambdify(self, recompile=False): """ Compile the function for discrete evaluation. """ if self._discrete_func is None or recompile == True: npfunc = sy.lambdify(sorted(self.free_symbols), self, 'numpy') if self._compile_type == "jit": self._discrete_func = jit(nopython=True)(npfunc) elif self._compile_type == "vectorize": self._discrete_func = vectorize([], nopython=True)(npfunc) else: raise NotImplementedError() return self._discrete_func
def test_power_integer(self): """ Test 0 ** -1. Note 2 ** <big number> returns an undefined value (depending on the algorithm). """ dtype = np.int64 f = vectorize(["int64(int64, int64)"], nopython=True)(power) a = np.array([5, 0, 6], dtype=dtype) b = np.array([1, -1, 2], dtype=dtype) expected = np.array([5, -2**63, 36], dtype=dtype) with self.check_warnings([]): res = f(a, b) self.assertPreciseEqual(res, expected)
def test_type_inference(self): global vector_add vector_add = vectorize([ bool_(double, int_), double(double, double), float_(double, float_), ])(add) cfunc = jit(func) self.assertEqual(cfunc(np.dtype(np.float64), np.dtype('i')), int8[:]) self.assertEqual(cfunc(np.dtype(np.float64), np.dtype(np.float64)), double[:]) self.assertEqual(cfunc(np.dtype(np.float64), np.dtype(np.float32)), float_[:])
def test_cuda_vectorize_device_call(self): ufunc = vectorize([float32(float32, float32, float32)], target='cuda')( cu_ufunc) N = 100 X = np.array(np.random.sample(N), dtype=np.float32) Y = np.array(np.random.sample(N), dtype=np.float32) Z = np.array(np.random.sample(N), dtype=np.float32) + 0.1 out = ufunc(X, Y, Z) gold = (X ** Y) / Z self.assertTrue(np.allclose(out, gold))
def _get_numba_ufunc(expr): """Construct a numba ufunc from a blaze expression Parameters ---------- expr : blaze.expr.Expr Returns ------- f : function A numba vectorized function Examples -------- >>> from blaze import symbol >>> import numpy as np >>> s = symbol('s', 'float64') >>> t = symbol('t', 'float64') >>> x = np.array([1.0, 2.0, 3.0]) >>> y = np.array([2.0, 3.0, 4.0]) >>> f = get_numba_ufunc(s + t) >>> f(x, y) array([ 3., 5., 7.]) See Also -------- get_numba_type compute_signature """ if isinstance(expr, Broadcast): leaves = expr._scalars expr = expr._scalar_expr else: leaves = expr._leaves() # we may not have a Broadcast instance because arithmetic expressions can # be vectorized so we use getattr s, scope = funcstr(leaves, expr) scope = dict((k, numba.jit(nopython=True)(v) if callable(v) else v) for k, v in scope.items()) func = eval(s, scope) sig = compute_signature(expr) return numba.vectorize([sig], nopython=True)(func)
def test_vectorize_error_in_objectmode(self): # An exception raised inside an object mode @vectorized function # is printed out and ignored. ufunc = vectorize(['int32(int32)'], forceobj=True)(uerror) a = numpy.arange(4, dtype='int32') b = numpy.zeros_like(a) with support.captured_stderr() as err: ufunc(a, out=b) err = err.getvalue() if sys.version_info >= (3, 4): self.assertIn("Exception ignored in: 'object mode ufunc'", err) self.assertIn("MyException: I'm here", err) else: self.assertRegexpMatches(err, r"Exception [^\n]* in 'object mode ufunc' ignored") self.assertIn("I'm here", err) self.assertTrue(numpy.all(b == numpy.array([1, 2, 0, 4])))
def _test_template_4(self, target): sig = [int32(int32, int32), uint32(uint32, uint32), float32(float32, float32), float64(float64, float64)] basic_ufunc = vectorize(sig, target=target)(vector_add) np_ufunc = np.add def test(ty): data = np.linspace(0., 100., 500).astype(ty) result = basic_ufunc(data, data) gold = np_ufunc(data, data) np.testing.assert_allclose(gold, result) test(np.double) test(np.float32) test(np.int32) test(np.uint32)
def test_vectorize_identity(self): sig = 'int32(int32, int32)' for identity in self._supported_identities: ufunc = vectorize([sig], identity=identity)(add) expected = None if identity == 'reorderable' else identity self.assertEqual(ufunc.identity, expected) # Default value is None ufunc = vectorize([sig])(add) self.assertIs(ufunc.identity, None) # Invalid values with self.assertRaises(ValueError): vectorize([sig], identity='none')(add) with self.assertRaises(ValueError): vectorize([sig], identity=2)(add)
def __init__(self, expression, arguments, argument_dtypes, return_dtype, verbose=False): self.expression = expression self.arguments = arguments self.argument_dtypes = argument_dtypes self.return_dtype = return_dtype self.verbose = verbose import numba argument_dtypes_numba = [getattr(numba, argument_dtype.name) for argument_dtype in argument_dtypes] argstring = ", ".join(arguments) code = ''' from numpy import * def f({0}): return {1}'''.format(argstring, expression) if verbose: print('Generated code:\n' + code) scope = {} exec(code, scope) f = scope['f'] return_dtype_numba = getattr(numba, return_dtype.name) vectorizer = numba.vectorize([return_dtype_numba(*argument_dtypes_numba)]) self.f = vectorizer(f)
def get_evaluate_functions(self): """Return compiled functions with fixed network parameters""" if self._ev_fs and self._evaluate_sync: return self._ev_fs ev_fs = [] param_symbols = self.parameter_symbols param_values = self.parameter_values input_symbols = list(self.input_symbols) input_len = len(input_symbols) for f in self.total_output_formulas: expr = f.subs(zip(param_symbols, param_values)) func = sympy.lambdify(input_symbols, expr) signature = 'float64(%s)' % ','.join(repeat('float64',input_len)) ev_fs.append(numba.vectorize(signature)(func)) self._evaluate_sync = True self._ev_fs = ev_fs return ev_fs
def template_vectorize(self, target): # build basic native code ufunc sig = [int32(int32, int32), uint32(uint32, uint32), float32(float32, float32), float64(float64, float64)] basic_ufunc = vectorize(sig, target=target)(vector_add) # build python ufunc np_ufunc = np.add # test it out def test(ty): data = np.linspace(0., 100., 500).astype(ty) result = basic_ufunc(data, data) gold = np_ufunc(data, data) self.assertTrue(np.allclose(gold, result)) test(np.double) test(np.float32) test(np.int32) test(np.uint32)
def test_type_inference(self): """This is testing numpy ufunc dispatch machinery """ global vector_add vector_add = vectorize([ bool_(double, int_), double(double, double), float_(double, float_), ])(add) cfunc = jit(func) def numba_type_equal(a, b): self.assertEqual(a.dtype, b.dtype) self.assertEqual(a.ndim, b.ndim) numba_type_equal(cfunc(np.dtype(np.float64), np.dtype('i')), bool_[:]) numba_type_equal(cfunc(np.dtype(np.float64), np.dtype(np.float64)), double[:]) # This is because the double(double, double) matches first numba_type_equal(cfunc(np.dtype(np.float64), np.dtype(np.float32)), double[:])