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-cudapackage 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))) # TrueRemember
- 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.