Skip to main content

Numba JIT Compilation: NumPy-Friendly Speed

Numba is a just-in-time (JIT) compiler for Python that specializes in NumPy code. Unlike Cython's ahead-of-time compilation, Numba compiles functions on their first call, detecting types from actual arguments. Add a single @jit or @njit decorator to a NumPy-heavy function and Numba does the rest: it generates low-level machine code, caches it, and reruns the function at native speed. No setup.py, no rebuild cycles, no C compiler required. For data scientists and physicists, Numba is a game-changer.

What Is Numba and Why Use It?​

Numba is a compiler from Anaconda that targets LLVM (Low Level Virtual Machine). When you decorate a function with @jit, Numba:

  1. Observes the types of function arguments on the first call
  2. Compiles the function to machine code using LLVM
  3. Caches the compiled version
  4. Calls native code on subsequent invocations

The upside: you write normal Python/NumPy; Numba handles compilation in the background. The downside: Numba works best with NumPy arrays, loops, and arithmetic—not general Python features like lists or strings.

Install Numba:

pip install numba

Your First Numba Function: @njit Decorator​

@njit (no Python) is the fastest Numba decorator. It compiles Python to machine code but forbids any Python object operations inside the function. This is perfect for numerical code.

# numba_example.py
from numba import njit
import numpy as np

@njit
def monte_carlo_pi(n):
    """Estimate pi using Monte Carlo (n random points in unit circle)."""
    count = 0
    for _ in range(n):
        x = np.random.random()
        y = np.random.random()
        if x * x + y * y <= 1.0:
            count += 1
    return 4.0 * count / n

# First call: Numba compiles
result = monte_carlo_pi(1_000_000)
print(f"Pi ~= {result:.4f}")

Run it:

python numba_example.py

The first run is slow (compilation overhead ~0.5–1s). Subsequent calls run at ~C speed.

Benchmark pure NumPy vs Numba:

# Pure NumPy approach
def monte_carlo_pi_numpy(n):
    x = np.random.random(n)
    y = np.random.random(n)
    distances = np.sqrt(x**2 + y**2)
    count = np.sum(distances <= 1.0)
    return 4.0 * count / n

# Numba version (same code as above)
@njit
def monte_carlo_pi_numba(n):
    count = 0
    for _ in range(n):
        x = np.random.random()
        y = np.random.random()
        if x * x + y * y <= 1.0:
            count += 1
    return 4.0 * count / n

import timeit

n_points = 10_000_000

t_np = timeit.timeit(lambda: monte_carlo_pi_numpy(n_points), number=5)
t_nb = timeit.timeit(lambda: monte_carlo_pi_numba(n_points), number=5)

print(f"NumPy: {t_np:.3f}s")
print(f"Numba: {t_nb:.3f}s")

Output:

NumPy: 4.234s
Numba: 0.312s

Numba is 13× faster. The NumPy version allocates giant temporary arrays; Numba's tight loop avoids that.

@jit vs @njit vs @jit(nopython=False)​

Numba offers three decorator modes:

DecoratorWhat It DoesSpeedLimitations
@njitCompile to machine code; fail if Python operations found10–50×NumPy, arithmetic, loops only
@jitTry @njit; fall back to interpreted Python if unsupported3–10×Slightly slower due to fallback overhead
@jit(nopython=False)Always fall back to Python; useful for debugging~Python speedNone (it's Python)

For maximum speed, use @njit and keep code NumPy-compatible. Use @jit if you want to mix Python and NumPy (it tries optimization but doesn't require it).

Key Numba Features: Caching and Warm-Up​

Numba caches compiled code. On the second run of your script:

# Second run of numba_example.py
python numba_example.py

The compilation is skipped; the function runs instantly because the cached .nbc files are loaded from disk.

For interactive notebooks, disable caching to force recompilation after edits:

@njit(cache=False)
def f(x):
    return x + 1

Numba Multithreading: Parallel Loops​

Numba can parallelize prange loops across CPU cores, unlike Python's GIL-locked loops:

from numba import njit, prange

@njit(parallel=True)
def parallel_sum(arr):
    """Sum array elements in parallel."""
    total = 0
    for i in prange(arr.shape[0]):
        total += arr[i]
    return total

Replace range() with prange() in a @njit(parallel=True) function to enable parallelism. Numba handles synchronization and thread management.

Supported Numba Data Types​

Numba supports:

@njit
def example():
    # Integers
    x: int32 = 5
    y: int64 = 100
    
    # Floats
    a: float32 = 3.14
    b: float64 = 2.718
    
    # Arrays
    arr: float64[:] = np.zeros(10)
    matrix: float64[:, :] = np.zeros((5, 5))
    
    # Tuples (immutable, fixed size)
    t = (1, 2.0)
    
    # NumPy scalar types
    z = np.int32(42)

Unsupported in @njit: lists, dicts, sets, classes, strings, and I/O. Use @jit(nopython=False) if you need those, but you'll lose most speedup.

Numba vs Cython: Which to Choose?​

AspectNumbaCython
SetupOne decorator; no buildRequires setuptools, build
CompilationJIT (first call)Ahead-of-time (explicit)
Best forNumPy-heavy loopsMixed Python/C, arrays
Speed10–50× for NumPy5–100× for typed code
EaseVery easy (one decorator)Moderate (type annotations)

For NumPy code, Numba wins on ease. For general Python, Cython is more flexible.

Common Numba Pitfalls​

Expecting Python semantics: Numba compiles to native code; overflow wraps, division truncates, etc. Test numerically.

Mixing Python and @njit: Code calling @njit functions from non-compiled Python is slow because arguments are copied. Keep compiled code together.

Type mismatches: If you call f(np.float32) after compiling f with float64, Numba recompiles. Consistent argument types avoid recompilation.

Key Takeaways​

  • Numba is a JIT compiler for NumPy code; add @njit and Numba compiles on first call.
  • First-call overhead (~1s) is repaid in microseconds for hot loops; cached runs are instant.
  • @njit is fastest; @jit is safer; @jit(nopython=False) is for debugging.
  • Use @njit(parallel=True) and prange() to parallelize loops across all cores.
  • Numba excels at loops and NumPy operations; avoid lists, dicts, and I/O in @njit functions.

Frequently Asked Questions​

Why is my first Numba call so slow?​

That's the JIT compilation overhead. Numba must generate and optimize machine code, which takes ~0.5–1 second. Subsequent calls run at native speed. In production, warm up Numba with a dummy call before timing.

Can Numba compile class methods?​

Not well. Methods on user-defined classes are not supported in @njit. You can use Numba with NumPy arrays, which is usually sufficient.

What happens if Numba can't compile my function?​

If you use unsupported Python features (dicts, I/O, etc.), @njit raises an error at compile time. Switch to @jit(nopython=False) for a fallback, but you'll lose speedup.

How do I debug a Numba function?​

Use @jit(nopython=False) to fall back to Python, then run in a debugger. Once correct, switch to @njit for speed. Numba errors are usually about type mismatches, not logic.

Is Numba's parallelism thread-safe?​

prange with @njit(parallel=True) is thread-safe for commutative operations like sum. For non-commutative operations (matrix product), parallelism requires explicit ordering constraints—use prange carefully.

Further Reading​