@vectorize and @guvectorize ufuncs

Numba · Practical · 12 min

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.

Check yourself

  1. What does a @guvectorize function with layout '(n)->()' applied to a (5, 3) array return?

Sign in to answer, see the explanations, and get this lesson back as a quick review later.

Official docs: numba.readthedocs.io/en/stable/user/vectorize.html