NumPy Support in Numba

Numba is designed to be used with NumPy and complement its capabilities. Numba supports:

  • Passing NumPy arrays as arguments, including structured dtypes
  • Creating compiled ufuncs and generalized ufuncs
  • Using a large subset of NumPy functions in nopython mode
In [ ]:
import numpy as np
import numba
from numba import jit

Numba Specialization by Dtype

Numba automatically uses multiple dispatch on compiled functions to allow different specialized implementations of the same function to be used. Suppose we have a function that clamps values to zero if they are below a particular magnitude:

In [ ]:
@jit(nopython=True)
def zero_clamp(x, threshold):
    # assume 1D array.  See later in this notebook for more general function
    out = np.empty_like(x)
    for i in range(out.shape[0]):
        if np.abs(x[i]) > threshold:
            out[i] = x[i]
        else:
            out[i] = 0
    return out        
In [ ]:
a_small = np.linspace(0, 1, 50)
zero_clamp(a_small, 0.3)

Now let's benchmark some different kinds of array inputs. We'll try:

  • int64
  • float32
  • float32 with a stride (elements not contiguous in memory)
In [ ]:
n = 10000
a_int16 = np.arange(n).astype(np.int16)
a_float32 = np.linspace(0, 1, n, dtype=np.float32)
a_float32_strided = np.linspace(0, 1, 2*n, dtype=np.float32)[::2]  # view of every other element
In [ ]:
%timeit zero_clamp(a_int16, 1600)
%timeit zero_clamp(a_float32, 0.3)
%timeit zero_clamp(a_float32_strided, 0.3)

We see different performance characteristics for each of these cases, even though they have the same number of input elements. Numba generated different machine code for each situation, which we can see if we look at the .signatures attribute of the compiled function:

In [ ]:
zero_clamp.signatures

When printed as strings, Numba array types have the form: array(dtype, dimensions, layout). The first signature therefore corresponds to a 1D array of float64 with C style layout (row-major order, no gaps between elements). The next two signatures are similar, but for int16 and float32 arrays. The final signature indicates an "any" layout array, which usually happens when you slice an array, and it no longer has a C or FORTRAN memory layout.

We can compare to a pure NumPy implementation and see the speed improvement that Numba has achieved through a combination of specialization and elimination of temporary arrays:

In [ ]:
def np_zero_clamp(x, threshold):
    return np.where(np.abs(x) > threshold, x, 0)
In [ ]:
%timeit np_zero_clamp(a_int16, 1600)
%timeit np_zero_clamp(a_float32, 0.3)
%timeit np_zero_clamp(a_float32_strided, 0.3)

Creating Ufuncs

Universal functions, typically called "ufuncs" for short, are functions that broadcast an elementwise operation across input arrays of varying numbers of dimensions. Most NumPy functions are ufuncs, and Numba makes it easy to compile custom ufuncs using the @vectorize decorator.

In [ ]:
from numba import vectorize
In [ ]:
@vectorize(nopython=True)
def ufunc_zero_clamp(x, threshold):
    if np.abs(x) > threshold:
        return x
    else:
        return 0
In [ ]:
%timeit ufunc_zero_clamp(a_int16, 1600)
%timeit ufunc_zero_clamp(a_float32, 0.3)
%timeit ufunc_zero_clamp(a_float32_strided, 0.3)

Note that for this simple ufunc, Numba is not as fast as the function with the manual looping, and in some cases, is the same speed as the example that called NumPy directly. This is not surprising as this function is very simple, and NumPy also uses compiled ufuncs. Numba @vectorize is generally most effective when creating ufuncs that are not a simple combination of existing NumPy operations.

Calling NumPy functions

Numba supports many, but not all, NumPy functions. Some functions also have limitations that prevent the use of some of the optional arguments in nopython mode. A full description can be found in the Supported NumPy Features page in the Numba Reference Manual.

Note that when using NumPy functions on arrays, Numba will also compile and optimize array expressions:

In [ ]:
def numpy_mpe(x, true):
    return (((x - true)/true)**2).mean()

numba_mpe = jit(nopython=True)(numpy_mpe)  # using jit as a function rather than a decorator

We can confirm both versions give the same answer:

In [ ]:
true_x = 0.1
x = np.random.normal(true_x, 1, size=100000)
numpy_mpe(x, true=true_x), numba_mpe(x, true=true_x)

And see the Numba version is faster:

In [ ]:
%timeit numpy_mpe(x, true=0.1)
%timeit numba_mpe(x, true=0.1)

If the scipy package is installed, Numba will also automatically make use of the optimized BLAS/LAPACK implementation that SciPy was compiled with. In the case of Anaconda, this is Intel MKL, but OpenBLAS is also common for builds of scipy. (Note that Numba is not itself compiled and linked against any BLAS implementation.) Most functions in numpy.linalg will be accelerated this way, as well as numpy.dot.

Numba will not run any faster than NumPy for individual linear algebra routines (since both translate to calls to the same underlying library), but you are able to use linear algebra calls inside your Numba-compiled functions without any loss of performance.