Numba: JIT Compilation for NumPy

Just-In-Time (JIT) compilation translates a function from bytecode to native machine code at the moment it is first called, rather than ahead of time. The compiled version is cached and reused on subsequent calls.

import numba

@numba.jit(nopython=True)
def sum_squares(arr):
    total = 0.0
    for i in range(len(arr)):
        total += arr[i] ** 2
    return total

The nopython=True flag (equivalently @numba.njit) forces Numba to compile entirely without the Python interpreter. If Numba cannot infer types, it raises an error rather than silently falling back.

Numba infers types from the arguments at the first call and generates machine code specialized for those types. You can also provide explicit type signatures:

@numba.njit('float64(float64[:])')
def norm_squared(x):
    s = 0.0
    for i in range(len(x)):
        s += x[i] * x[i]
    return s

When a signature is provided, Numba compiles eagerly (at decoration time) rather than lazily (at first call).

@numba.vectorize creates a NumPy universal function (ufunc) from a scalar kernel. The resulting function automatically broadcasts over arrays and supports output dtypes, reduction, and accumulation.

@numba.vectorize(['float64(float64, float64)'])
def clip_add(x, y):
    result = x + y
    if result > 1.0:
        return 1.0
    elif result < -1.0:
        return -1.0
    return result

This compiles the scalar function and wraps it as a ufunc that operates element-wise on arrays of any shape.

@numba.cuda.jit compiles a Python function into a CUDA GPU kernel. The programmer must specify the grid and block dimensions explicitly:

from numba import cuda

@cuda.jit
def vector_add(a, b, out):
    idx = cuda.grid(1)
    if idx < out.size:
        out[idx] = a[idx] + b[idx]

threads_per_block = 256
blocks = (N + threads_per_block - 1) // threads_per_block
vector_add[blocks, threads_per_block](d_a, d_b, d_out)

Memory must be explicitly transferred between host and device using cuda.to_device() and d_arr.copy_to_host().

Numba offers two compilation modes:

• nopython mode (@njit or nopython=True): Compiles the entire function to machine code. No Python objects allowed; all types must be inferable. This is the fast path.

• object mode (nopython=False, the old default): Falls back to Python objects when type inference fails. Provides minimal speedup and should be avoided.

Additional options include cache=True (saves compiled code to disk), parallel=True (auto-parallelizes eligible loops), and fastmath=True (relaxes IEEE 754 for extra speed).

When parallel=True is set, Numba can distribute loop iterations across CPU cores. Use numba.prange instead of range to mark parallelizable loops:

@numba.njit(parallel=True)
def parallel_sum(arr):
    total = 0.0
    for i in numba.prange(len(arr)):
        total += arr[i] ** 2
    return total

Numba automatically handles thread creation, synchronization, and reduction operations (sum, min, max).