Numba JIT Compilation

Basic Usage

from numba import jit, njit, prange, int64, float64
import numpy as np

@njit  # nopython mode (recommended)
def sum_array(arr):
    return np.sum(arr)

@njit(float64(float64[:], int64))  # Type signature
def get_element(arr, idx):
    return arr[idx]

@njit(cache=True)  # Cache to disk
def cached_function(x):
    return x * 2

@njit(parallel=True)  # Parallel execution
def parallel_sum(arr):
    total = 0.0
    for i in prange(arr.shape[0]):
        total += arr[i]
    return total

Common Patterns

@njit
def process_matrix(matrix):
    """2D array processing"""
    return matrix * 2

@njit
def compute_stats(arr):
    """Multiple return values"""
    return arr.mean(), arr.std()

@njit
def pairwise_distances(X):
    """Pairwise Euclidean distances"""
    n = X.shape[0]
    distances = np.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            distances[i, j] = distances[j, i] = np.sqrt(np.sum((X[i] - X[j])**2))
    return distances

NumPy Support and Limitations

Supported: np.zeros, np.ones, np.sum, np.mean, np.std, np.min, np.max, np.sqrt, np.exp, np.log, np.dot, np.outer, np.linalg.norm

Limitations: No dynamic typing in nopython mode, limited stdlib, use @jitclass for classes

Debug: Set NUMBA_DISABLE_JIT=1 or config.DISABLE_JIT = True