numba

name

numba

description

Use this skill whenever working with Numba JIT compilation, @nb.njit, @nb.jit, parallel=True/False, prange, numba.typed, numba.cuda, numba.vectorize, numba.guvectorize, or debugging Numba compilation errors. Trigger on: slow numpy loops being accelerated, nopython mode errors, Numba type inference failures, reflected list warnings, cache=True, fastmath, nogil, Numba + NumPy interop, or any task involving @njit on scientific/numerical code.

Numba: JIT Compilation for Scientific Python

What Numba does

Numba compiles Python functions to native machine code at call time using LLVM. It targets numerical code operating on NumPy arrays and Python scalars. The compiled function runs at C/Fortran speed with no interpreter overhead.

Key constraint: only a subset of Python is supported inside JIT-compiled functions. The supported subset is called nopython mode (NPM). When in doubt, always target NPM.

---

Decorator reference

import numba as nb
import numpy as np

@nb.njit — the main decorator

@nb.njit
def f(x, y):
    return x + y

Equivalent to @nb.jit(nopython=True). Always prefer @nb.njit over @nb.jit. @nb.jit without nopython=True silently falls back to object mode on failure — this masks bugs and gives no speedup.

Common options:

@nb.njit(
    parallel=False,   # True: enable automatic parallelisation + prange
    cache=True,       # persist compiled bitcode to __pycache__, avoid recompile
    fastmath=True,    # allow reassociation, FMA, no-NaN, no-INF (UNSAFE for QMC)
    nogil=True,       # release the GIL: allows Python threading around nb calls
    boundscheck=False # default False; set True only during debugging
)
def f(x): ...

⚠️ fastmath=True is dangerous for QMC/quantum chemistry: it allows floating-point reassociation that breaks energy conservation. Use only for non-critical kernels (geometry, index arithmetic).

@nb.njit(parallel=True) + nb.prange

@nb.njit(parallel=True)
def sum_rows(A):
    result = np.zeros(A.shape[0])
    for i in nb.prange(A.shape[0]):   # prange → parallel loop
        for j in range(A.shape[1]):
            result[i] += A[i, j]
    return result

nb.prange replaces range in the outermost loop only. Inner loops stay serial. Only use parallel=True when iterations are independent (no loop-carried dependencies).

In pycasino, all kernels use parallel=False. Parallelism is achieved at the MPI level (independent walkers). Do not add parallel=True without profiling — the threading overhead can hurt for small arrays.

@nb.vectorize — ufunc factory

@nb.vectorize(['float64(float64, float64)'], nopython=True)
def clip(x, lo, hi):
    if x < lo: return lo
    if x > hi: return hi
    return x

Produces a NumPy ufunc that broadcasts automatically. Signature list is optional but speeds up dispatch. Without it, Numba infers at first call.

@nb.guvectorize — generalised ufunc

For functions that operate on arrays of fixed rank (e.g., matrix × vector):

@nb.guvectorize(['void(float64[:,:], float64[:], float64[:])'],
                '(m,n),(n)->(m)', nopython=True)
def matvec(A, x, out):
    for i in range(A.shape[0]):
        out[i] = 0.0
        for j in range(A.shape[1]):
            out[i] += A[i, j] * x[j]

---

Supported Python/NumPy inside @njit

Works ✓

• Arithmetic, comparisons, boolean logic
• if/elif/else, for, while, break, continue, return
• Tuples (fixed-length), basic unpacking
• range(), len(), abs(), min(), max(), round()
• NumPy array creation: np.zeros, np.ones, np.empty, np.full
• NumPy operations: np.sum, np.dot, np.cross, np.sqrt, np.exp, np.log
• np.linalg.norm, np.linalg.det, np.linalg.inv, np.linalg.solve
• np.linalg.eigh, np.linalg.svd (limited)
• Array indexing, slicing, boolean indexing (basic)
• math.pi, math.e, math.sin, math.cos, math.exp, math.sqrt
• numba.typed.List, numba.typed.Dict
• Calling other @njit functions (inlined automatically)

Does NOT work ✗

• list, dict, set (Python built-in containers — use numba.typed.* or arrays)
• print with f-strings (use print(x) with scalars only)
• try/except
• Generator expressions, list comprehensions (sometimes work, often don't)
• isinstance, hasattr, getattr
• Classes (except @nb.experimental.jitclass)
• **kwargs in JIT-compiled functions
• scipy.* (not supported inside njit)
• np.einsum — NOT supported in nopython mode; unroll manually or use dot/matmul

---

Type system and type inference

Numba infers types from the first call arguments. The compiled function is specialised per type signature. Calling with different dtypes creates multiple compiled versions.

Explicit signatures (optional, locks types, speeds first call)

@nb.njit('float64[:](float64[:], float64)')
def scale(arr, factor):
    return arr * factor

Signature syntax: 'return_type(arg1_type, arg2_type)'

Common type strings:

Type string	Python/NumPy meaning
float64	np.float64 scalar
float32	np.float32 scalar
int64	np.int64 scalar
boolean	bool
float64[:]	1-D C-contiguous float64 array
float64[:,:]	2-D C-contiguous float64 array
float64[::1]	1-D array, explicitly C-contiguous
float64[:, ::1]	2-D array, C-contiguous (last dim contiguous)
UniTuple(float64, 3)	tuple of 3 float64

Contiguity matters

Numba generates fastest code for C-contiguous arrays. If you slice a non-contiguous view, pass np.ascontiguousarray(x) before calling the JIT function.

---

Debugging Numba errors

Step 1: read the type error

numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
  - argument 0: cannot determine Numba type of <class 'list'>

→ A Python list was passed. Replace with np.array(...) or numba.typed.List.

Step 2: use nb.njit with cache=False and inspect

@nb.njit(cache=False)
def f(x): ...

f.inspect_types()        # shows inferred types for each variable
f.inspect_llvm()         # shows generated LLVM IR
f.inspect_asm()          # shows native assembly

Step 3: common mistakes

Reflected list warning:

NumbaTypeSafetyWarning: unsafe cast from int64 to int32 (reflected list)

→ Numba is boxing/unboxing a Python list at the boundary. Move the list inside the function or convert to np.array before calling.

Object mode fallback (silent, @jit only):

NumbaWarning: Function "f" was compiled in object mode without forceobj=True

→ Switch to @nb.njit to make this a hard error. Find and fix the unsupported construct.

cannot unify float64 and int64: Numba sees two branches returning different types. Make types consistent:

# Bad
if cond:
    return 0       # int
else:
    return 0.0     # float
# Good
    return 0.0
    return 0.0

---

Memory and array patterns

Output arrays: allocate outside, pass in

@nb.njit
def fill(out, x):
    for i in range(out.shape[0]):
        out[i] = x * i

out = np.empty(N, dtype=np.float64)
fill(out, 3.14)

Avoids allocation inside hot loops.

Returning arrays from @njit

Numba can return arrays allocated inside the function:

@nb.njit
def make(n):
    out = np.empty(n, dtype=np.float64)
    ...
    return out

This is fine. The array is heap-allocated by Numba and ownership is transferred to Python.

Slices and views

Array slices inside @njit are views (no copy). Safe to pass slices as arguments:

@nb.njit
def process(row):   # row is a 1-D view
    ...

for i in range(A.shape[0]):
    process(A[i, :])   # no copy, passes a view

---

Compilation cache

@nb.njit(cache=True)
def f(x): ...

Compiled bitcode is stored in __pycache__/. On the next run Numba skips compilation if the source has not changed. Essential for scripts that import many JIT functions.

Invalidation: cache is invalidated when the function source changes, the Numba version changes, or the NumPy version changes. Delete __pycache__ manually if you see stale-cache bugs.

---

Numba + pycasino patterns

The parallel=False convention

All pycasino kernels use parallel=False. This is intentional: each MPI rank occupies one core, and thread-level parallelism would conflict. Do not change this without benchmarking.

Calling @njit from Python (warm-up)

First call triggers compilation (can take seconds). For benchmarking, always call once with representative input before timing:

# Warm up
f(np.zeros(10, dtype=np.float64))
# Now benchmark
t0 = time.perf_counter()
f(real_input)
print(time.perf_counter() - t0)

Or use nb.njit(cache=True) to amortise the cost across runs.

Passing structured data

Numba does not support Python dataclasses or arbitrary objects. The pycasino pattern is to extract arrays from objects before calling JIT kernels:

# Python side (outside njit)
def value(self, r_e):
    return _value_jit(r_e, self.coefficients, self.exponents, self.shell_map)

@nb.njit
def _value_jit(r_e, coefficients, exponents, shell_map):
    ...

This keeps class logic in Python and numerical hot loops in Numba.

np.linalg inside @njit

Supported functions: det, slogdet, inv, solve, norm, eigh, cholesky. Not supported: lstsq, matrix_rank, pinv.

For pinv, implement via SVD:

@nb.njit
def pinv(A):
    U, s, Vt = np.linalg.svd(A)
    tol = 1e-12 * s[0]
    s_inv = np.where(s > tol, 1.0 / s, 0.0)
    return (Vt.T * s_inv) @ U.T

---

Performance checklist

1. Profile first — use cProfile or line_profiler to confirm Numba functions are the bottleneck before optimising.
2. Check dtypes — float32 is 2× faster on GPU/SIMD but loses precision. For QMC always use float64.
3. Avoid Python objects inside hot loops — every Python object lookup breaks the JIT. Keep everything as scalars or numpy arrays.
4. Prefer 1-D loops over numpy broadcasting inside @njit — Numba vectorises explicit loops well; broadcasting sometimes defeats optimisation.
5. Use nb.prange only at the outermost loop — nesting prange inside prange is not supported.
6. cache=True for all stable functions — eliminates recompilation overhead in production runs.
7. Avoid np.einsum — unroll with explicit loops or use @ (matmul operator).

---

Useful environment variables

NUMBA_NUM_THREADS=4      # number of threads for parallel=True kernels
NUMBA_CACHE_DIR=/tmp/nb  # redirect cache away from source tree
NUMBA_DISABLE_JIT=1      # disable all JIT (pure Python fallback, for debugging)
NUMBA_DEBUG_TYPEINFER=1  # verbose type inference log
NUMBA_DUMP_IR=1          # dump LLVM IR for compiled functions

NUMBA_DISABLE_JIT=1 is invaluable for debugging: all @njit functions run as plain Python, so pdb, print, and traceback work normally.