Writing inside the supported subset from the start avoids fighting compile errors later.
Numba compiles a subset of Python and NumPy. Inside an @njit function you can use:
- numbers (
int,float,complex,bool),None, tuples and named tuples, if,for,while,break,continue,range,enumerate,zip,- NumPy arrays of numeric, boolean and some datetime dtypes, with indexing, slicing, views and most arithmetic,
- a large part of the NumPy API: creation functions (
np.zeros,np.empty,np.arange), math ufuncs, reductions (sum,mean,argmax),np.sort,np.linalg(via SciPy's LAPACK) and thenp.randomlegacy functions, - the
mathandcmathmodules, strings with basic operations, lists and dicts (typed, homogeneous), sets, simple classes via@jitclass, and calls to other@njitfunctions.
What does not work: pandas, most of the standard library, arbitrary Python objects, try with arbitrary exception types beyond basics, generators that escape, dicts or lists mixing types, and many keyword arguments of NumPy functions (for example dtype= or axis= support varies by function).
Containers must be homogeneous: a list of floats is fine; [1, "a"] is a typing error. Arrays are the preferred data structure, because they map directly to raw memory.
When something is unsupported, the usual fix is to restructure: do I/O and pandas work outside, then pass plain NumPy arrays into a small compiled kernel. The reference pages list supported Python features and NumPy functions in detail.
Example
import math
import numpy as np
from numba import njit
@njit
def stats(a):
out = np.empty(3)
out[0] = a.mean()
out[1] = math.sqrt(((a - a.mean()) ** 2).sum() / a.size)
out[2] = np.argmax(a)
return out
print(stats(np.array([1.0, 2.0, 3.0, 10.0]))) # [4. 3.53553391 3. ]
@njit
def bad(xs):
return [x for x in xs] + ["oops"] # mixes float and str
try:
bad(np.ones(2))
except Exception as e:
print(type(e).__name__) # TypingErrorRemember
- Numbers, tuples, loops, NumPy arrays and much of the NumPy API compile
- Containers must hold a single type
- Keep pandas, I/O and objects outside; pass arrays into compiled kernels
Watch out for
- Assuming every keyword argument of a NumPy function is supported inside @njit.