What Numba does: @njit

Numba · Essentials · 9 min

It is the cheapest way to make a slow Python loop over numbers run at C speed without leaving Python.

Pure Python loops are slow because every operation goes through the interpreter and works on boxed objects. Numba is a just-in-time (JIT) compiler: when you call a decorated function, it reads the function's bytecode, infers concrete types for every variable from the argument types, and uses LLVM to generate native machine code. Later calls with the same argument types jump straight into that compiled code.

The decorator you will use most is @njit, short for @jit(nopython=True). In nopython mode the whole function is compiled without touching the Python interpreter, which is where the speed comes from. Since Numba 0.59, plain @jit also defaults to nopython mode, so the two behave the same; @njit is still the common spelling.

Numba works best on code that:

  • loops over NumPy arrays or ranges of numbers,
  • uses scalars, tuples, NumPy arrays and a subset of Python and NumPy functions,
  • would be awkward or memory-hungry to express as vectorized NumPy.

It does not speed up code that calls pandas, does string-heavy work or uses arbitrary Python objects; such code fails to compile in nopython mode with a TypingError.

The function stays a normal Python callable: call it with arrays and numbers as usual. The original pure Python version is available as func.py_func, which is handy for testing that both agree.

Example

import numpy as np
from numba import njit

@njit
def sum_of_squares(a):
    total = 0.0
    for x in a:           # a plain loop, compiled to machine code
        total += x * x
    return total

a = np.arange(1_000_000, dtype=np.float64)
print(sum_of_squares(a))           # 3.3333283333312755e+17
print(sum_of_squares.py_func(a[:4]))  # 14.0 (the original Python function)
print(sum_of_squares(a[:4]))       # 14.0

Remember

  • @njit compiles a function to native code on its first call
  • Nopython mode means no interpreter in the loop: that is where speed comes from
  • func.py_func is the original Python function

Watch out for

  • Decorating code that uses pandas or arbitrary objects and expecting it to compile.

Check yourself

  1. Why does a Numba-compiled loop run much faster than the same pure Python loop?

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/5minguide.html