Numba vs NumPy vs others

Numba · Practical · 9 min

Adding @njit everywhere wastes compile time; knowing where it wins focuses the effort.

Vectorized NumPy already runs compiled C loops, so wrapping a single NumPy call such as a.sum() or a @ b in @njit gives no speedup: Numba calls comparable code, and matrix products go to the same BLAS. Numba wins when NumPy would need multiple passes or temporaries, or when the algorithm is inherently sequential:

  • Expressions creating many temporary arrays, like np.sqrt(x**2 + y**2) * np.exp(-z), which allocate an array per operation. A Numba loop fuses them into one pass with no temporaries.
  • Recurrences, state machines, early exits, and searches.
  • Algorithms with irregular access patterns: graphs, histograms with custom binning, particle simulations, dynamic programming.
  • Code that would need np.vectorize (which is just a Python loop) or Python-level loops in NumPy.

Numba does not help with I/O-bound code, pandas and string processing, or code dominated by calls to large NumPy or BLAS routines.

Alternatives to know:

  • Cython: ahead-of-time compiled, good for wrapping C libraries and shipping wheels, more setup.
  • JAX: array programming with automatic differentiation and GPU/TPU support, functional style.
  • NumExpr: fuses array expressions without writing loops.
  • For GPUs, Numba's CUDA support now lives in the separate numba-cuda package maintained by NVIDIA.

Measure first: profile, find the hot function, then decide.

Example

import time
import numpy as np
from numba import njit

x = np.random.default_rng(1).random(5_000_000)

def numpy_version(x):
    return np.sqrt(x**2 + 1.0) * np.exp(-x)   # several temporary arrays

@njit
def numba_version(x):
    out = np.empty_like(x)
    for i in range(x.size):                    # one fused pass, no temporaries
        out[i] = np.sqrt(x[i] ** 2 + 1.0) * np.exp(-x[i])
    return out

numba_version(x[:10])                          # warm up
for f in (numpy_version, numba_version):
    t = time.perf_counter(); f(x)
    print(f.__name__, f"{time.perf_counter() - t:.3f}s")
print(np.allclose(numpy_version(x), numba_version(x)))  # True

Remember

  • Single NumPy calls are already compiled; wrapping them in @njit gains nothing
  • Numba wins on fused expressions, recurrences and irregular algorithms
  • Profile first; consider Cython, JAX or NumExpr where they fit better

Watch out for

  • Expecting @njit around a large matrix multiplication to beat NumPy's BLAS.

Check yourself

  1. Which function is most likely to get a big speedup from @njit?

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/performance-tips.html