Custom ufuncs plug into NumPy broadcasting and reductions while running compiled code.
@vectorize takes a function written for scalars and builds a true NumPy ufunc from it. The ufunc then works on scalars and arrays of any shape, with full broadcasting, out= arguments and ufunc methods like .reduce and .accumulate.
@vectorize(["float64(float64, float64)"])
def rel_diff(x, y): ...With a list of signatures, the ufunc is compiled eagerly for those types. Without signatures you get a dynamic ufunc that compiles new loops as it meets new types. The target option selects "cpu" (default, single thread) or "parallel" (multithreaded, worth it only for large arrays).
@guvectorize builds a generalized ufunc that operates on sub-arrays rather than single elements. You declare a layout string like "(n),()->(n)", and the function receives array slices and writes results into an output argument (it does not return). NumPy loops over the remaining leading dimensions for you, so a function written for one row automatically applies to every row of a 2D array.
Compared with @njit, ufuncs give you NumPy-native behaviour for free: broadcasting rules, dtype dispatch and integration with code that expects ufuncs. Use @njit when you need control flow across elements, such as recurrences, which a ufunc cannot express.
Example
import numpy as np
from numba import vectorize, guvectorize, float64
@vectorize([float64(float64, float64)])
def rel_diff(x, y):
return 2 * (x - y) / (x + y)
a = np.array([1.0, 2.0, 4.0])
print(rel_diff(a, 2.0)) # [-0.66666667 0. 0.66666667] (broadcast)
print(rel_diff(a[:, None], a).shape) # (3, 3)
@guvectorize([(float64[:], float64, float64[:])], "(n),()->(n)")
def clip_row(row, limit, out):
for i in range(row.shape[0]):
out[i] = min(row[i], limit)
m = np.array([[1.0, 5.0], [7.0, 2.0]])
print(clip_row(m, 3.0)) # [[1. 3.]
# [3. 2.]]Remember
- @vectorize turns a scalar function into a broadcasting NumPy ufunc
- @guvectorize works on sub-arrays described by a layout string and writes to an output argument
- Use @njit when elements depend on each other
Watch out for
- Returning a value from a @guvectorize function instead of filling the output argument.