3 Numba Tricks for Python Runtime Optimization: Compile the Loop, Widen It, Stop Recompiling It

Python loops are slow for a reason that has nothing to do with the loop. Every iteration of a plain Python for loop sends the interpreter back to the same question: what types are these objects, and which operation should I dispatch? For a reduction over a ten-million-element NumPy array, that question gets asked ten million times. The arithmetic itself is trivial. The bookkeeping is not.

Numba exists to delete that bookkeeping. It compiles a numeric Python loop to machine code without leaving the Python environment, without a rewrite in C, and without forcing code into vectorized shapes it never wanted. You add a decorator, and on the first call Numba reads the argument types, compiles a specialization for them, and from then on runs native machine code.

The decorated version of a reduction differs from the plain one by exactly one line. That is the pitch, and it mostly holds. But when Numba code underperforms, the compiler is nearly never the culprit. The problem sits at the boundary around the compiled code, and it takes one of three forms: you did not cross it, you did not make it wide enough, or you are paying to cross it on every run.

Three decorators address those three failures, and each section below names the failure it fixes. All timings quoted here were checked against Numba 0.67.0 on a 4-core x86-64 laptop, with a 10,000,000-element float64 array.

Trick 1: Stay inside nopython mode

The failure this fixes: you did not cross the boundary.

@njit is shorthand for @jit(nopython=True), and the nopython part is the whole point. Nopython mode produces much faster code, but it has limitations. It compiles only the subset of Python and NumPy for which it can assign types to everything it sees. Step outside that subset and Numba either raises an error or, if you used plain @jit without the flag, quietly drops into object mode.

Object mode is where gains go to die. The function still runs, still returns the right answer, and still pays interpreter overhead on every operation it could not type. You get the ceremony of a compiled function with almost none of the speed. Using Numba well is mostly a matter of keeping the hot function inside the subset it can type, which in practice means numeric scalars, typed arrays, and straightforward control flow.

The practical rule follows from that. If a function needs to call back into arbitrary Python objects, do that work outside the compiled function and pass in plain arrays. Keep the hot loop boring. Boring is what compiles.

Trick 2: Parallelize with parallel=True and prange

The failure this fixes: you did not make the boundary wide enough.

A compiled function still runs on one core by default. Numba’s job is to make that one core fast, not to find more of them. Adding parallel=True to the decorator tells Numba to try to parallelize the function, and swapping range() for prange() tells it which loop to target. The loop body does not change.

from numba import njit, prange

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

That looks unsafe. A shared accumulator incremented from multiple threads is the textbook race condition, and in most languages this code would be wrong. Numba recognizes total += ... as a reduction. It splits the range across threads, gives each thread a private accumulator, and combines the partial results at the end.

On the test machine, the serial @njit reduction over ten million elements ran in roughly 8 ms; the parallel=True version with prange came in at roughly 2.5 ms, a speedup of about 3.2x on 4 cores. The source change is one keyword, and the gain scales with available cores up to the point where memory bandwidth becomes the limit rather than arithmetic.

One caveat on the example: the accumulator here is a scalar, which is the case Numba’s reduction recognition is built for. If you try the same pattern with a NumPy array accumulator (out += arr[i]), the recognition does not apply the same way and you can get wrong results or a compile error. Use a scalar accumulator, or restructure the loop so each iteration writes only to its own index.

The same reduction recognition applies to -=, *=, /=, max and min. If your loop is a sum, a product, a running maximum or a running minimum, prange plus parallel=True is usually a free win. If your loop writes to arr[i] from iteration i only, that is also safe. If it writes to arr[i - 1] or accumulates into a shared list, the recognition does not apply and you have to restructure before parallelizing.

Trick 3: Cache compiled results with cache=True

The failure this fixes: you are paying to cross the boundary on every run.

Compilation happens on the first call. That is fine once, and wasteful every time after, because a fresh process starts with an empty compilation cache. Whether that matters depends entirely on how often you start a process.

For a script run once a day, the first-call compile cost is negligible against the work it does. For a tool invoked twenty times an hour, that cost can be the majority of runtime. The function is fast; you are just paying the compiler’s bill over and over for identical input types.

cache=True writes the compiled result to disk beside the source file. A later run loads it instead of recompiling. Against the parallel @njit implementation on the test machine, the first call cost roughly 1.4 s to compile and the cached second process started in roughly 30 ms, a speedup of more than 40x on startup. That is a large relative gain and a small absolute one, which is the expected result: the parallel version had already made the compute cheap enough that the compile step was no longer the dominant term in a long-running workload. The trick pays off in proportion to how often you restart the process, not how heavy the loop is.

The decision rule is therefore about process lifetime, not loop size. If your code runs inside a long-lived service or a notebook kernel that stays up for hours, cache=True buys you almost nothing and you can skip it. If your code runs as a short-lived CLI, a cron job, a CI step or a serverless function that pays startup cost on every invocation, turn it on.

What caching quietly breaks

cache=True has failure modes worth knowing before you trust it in a pipeline.

A global variable that the function reads is frozen at its compile-time value. On a later run, the cached code loads with the old value baked in and does not rebind to whatever the global holds now. If a constant, a threshold or a configuration flag lives at module scope and feeds into a cached function, changing it will not change the result.

Cache invalidation also fails to recognize changes in symbols defined in a different file. Editing a helper function in another module can leave you running previously compiled code, because the cache key does not see across that boundary. The fix is to clear the cache directory when you change anything the function depends on, or to keep dependencies inside the same file.

Caching a parallel=True function has a rougher history than caching a plain one. Before relying on it, verify the cache is actually being hit rather than assuming it. Set NUMBA_DEBUG_CACHE=1 in the environment and watch the output, or check the __pycache__ directory for the generated files. A cache that silently misses is worse than no cache, because it hides the cost you thought you had removed.

Compile the loop, widen it, stop recompiling it

One function body, three decorators, one underlying question. Is the work inside the compiled boundary? That is Trick 1, and @njit answers it by refusing to compile what it cannot type. Is the whole machine inside it? That is Trick 2, and parallel=True with prange answers it by treating a recognized reduction as safe to split. Are you paying to cross it more than once? That is Trick 3, and cache=True answers it by persisting the compiled artifact to disk.

Each trick moves the boundary rather than improving the compiler. That is the useful mental model, because it tells you where to look when the numbers disappoint. A function that got slower after you decorated it is probably dropping into object mode. A parallel version that barely beats the serial one is probably not parallelizing the loop you thought, or is dominated by a reduction Numba did not recognize. A tool that feels slow on startup is probably recompiling on every invocation.

Numba’s compiler is good at what it does. When your code underperforms anyway, the boundary around the compiled code is the place to look.

For a related look at the same instinct in pure Python, where the language already provides the mechanism and the work is understanding what it guarantees, see 7 Advanced Python Tricks That Use What the Language Already Promises You.

For more on this, see token compression prompt optimization.

For more on this, see statsmodels series tricks stop.