numba-utils

Battle-tested building blocks for production Numba workloads.

✓ Zero dependencies beyond NumPy + Numba  ·  ✓ Callable inside @njit  ·  ✓ Honest benchmarks  ·  ✓ Diagnostics

                    Numba
                      ▲
                      │
                 numba-utils
      ┌───────────────┼───────────────┐
   Arrays        Collections       Parallel
   Algorithms    Random            Profiling
   Decorators    Graph             Diagnostics
   Stats         Testing

from numba import njit
from numba_utils import topk

@njit
def winners(scores):
    return topk(scores, 10)     # O(n), no full sort — runs in nopython mode

from numba import njit
from numba_utils import PriorityQueue, SparseSet

@njit
def simulate(n_events):
    events = PriorityQueue(n_events)    # constructed in nopython mode
    active = SparseSet(100_000)         # O(1) add/discard/contains/clear
    ...

from numba_utils import compare

compare(numpy_impl, njit_impl, args=(values,))
# 31x on a fused kernel — JIT compilation excluded automatically

And the part that almost no library ships — diagnostics for compiled code:

>>> from numba_utils import diagnostics
>>> diagnostics.check(fn)
⚠ cache=True may crash when loaded across processes (farms, network FS)
  → NUMBA_UTILS_CACHE=0 in the environment, before the first import
⚠ fastmath=True relaxes IEEE 754 — not for exact/reproducible results

Install

pip install numba-utils

Zero build steps: NumPy and Numba are the only dependencies, and Numba brings its own LLVM-based compiler.

Why this exists

After enough numerical projects — long-running Monte Carlo engines, solvers, simulation farms — you realize you've rewritten the same binary search, the same typed collections, the same sampling algorithms and the same benchmarking helpers again. And debugged the same Numba production surprises again.

numba-utils ships more than code. It ships the production knowledge that usually stays trapped inside numerical projects — as kernels with the pitfalls engineered around, as diagnostics, and as documentation. It does not compete with Numba: it builds on top of it.

Why not...

• heapq / collections? They can't be called from nopython mode. The containers here are jitclasses usable inside @njit.
• NumPy? Many helpers are built to run inside compiled kernels, where NumPy calls can't reach. Where NumPy is faster (bandwidth-bound sweeps, its SIMD sort), BENCHMARKS.md says so.
• SciPy? A heavy dependency that isn't njit-callable; this stays at NumPy + Numba and works in the compiled path.
• Numba itself? Numba is a compiler. numba-utils is a standard library on top of it.

Design principles

1. Performance First — nothing ships without a benchmarked justification.
2. No Hidden Magic — thin, readable layers over Numba; nothing rewrites your code.
3. Numba Compatible — everything callable from your own @njit code, no hacks.
4. Minimal APIs — topk(arr, 10), not twenty keyword parameters.
5. Benchmark Honesty — losses are published next to the wins.

The identity behind these: docs/philosophy.md.

Modules

Core — decorators, arrays, algorithms, stats (logsumexp, softmax, weighted_quantile) · Performance — parallel (complete operations, not prange wrappers; bit-exact chunked_reduce), profiling (JIT excluded by default), diagnostics · Data structures — collections (dtype-generic factories), graph (BFS/DFS, toposort, Dijkstra, UnionFind over CSR), random (including the stateless counter-based Philox RNG, bit-identical to np.random.Philox) · Developer tools — testing (reference validation + stochastic asserts), config

Full API: docs/modules.md · Runnable code: examples/

Benchmark honesty

Every algorithm states whether it is faster, similar but more ergonomic, or slower but solving a problem unavailable elsewhere. BENCHMARKS.md contains losing rows on purpose: they tell you when NOT to use a function. Backed in-repo by reproducible benchmarks/, 300+ reference-validated tests (why there's no coverage badge), and CI running all of it. Trade-off records: docs/design/.

Used in

Patterns extracted from real workloads: Monte Carlo equity engines, game-theory solvers (CFR), quantitative research, scientific simulations, optimization loops.

Works with cachau

Result caching composes cleanly on top of numba-utils: the decorator aliases return real Numba dispatchers, so cachau fingerprints them — including the options the aliases inject (fastmath, parallel) and global configure() / NUMBA_UTILS_* overrides. @cache goes outermost:

from numba_utils.decorators import njit_fast
from cachau import cache

@cache(persist=True, max_memory="2GB")
@njit_fast
def simulate(values, iterations):
    ...

A cachau HIT skips execution entirely; a MISS still benefits from cache=True's compiled-code cache. Flipping a semantic option (fastmath, parallel, dev-mode boundscheck) changes the cache identity, so stale results are never served across semantics. Verified by cachau's integration suite.

Status

All three roadmap phases are shipped (see ROADMAP.md); what comes next is user-driven — open an issue.

• Phase 1 (0.1.x) — decorators, profiling, diagnostics, arrays, algorithms, random, collections, parallel, testing; PyPI release
• Phase 2 (0.2.0) — dtype-generic collections, stable_argsort, lexsort, the graph/ module (BFS/DFS, toposort, Dijkstra, UnionFind over CSR) and the stats/ module (logsumexp, softmax, weighted_quantile)
• Phase 3 (0.3.0) — Monte Carlo primitives from the first real-user feedback: counter-based Philox RNG (bit-identical to np.random.Philox), partial Fisher–Yates sampling, combination_table, bit-exact serial/parallel chunked_reduce, stochastic assertions
• 0.3.1–0.3.3 — four rounds of adversarial audit absorbed, each fixed and released the same day it was reported: memory-safety hardening, contract corrections, ~150 attack vectors (CHANGELOG.md)
• Phase 4 (0.4.0) — contributions from a production CFR solver: disjoint_rank_aggregate (reach-weighted all-pairs comparison with exact set-disjointness by inclusion–exclusion), the certification primitives (mutation_screams, assert_within_se, the reach² guard), and documented parallel patterns (Hogwild, factorized aggregation)

Development

python -m venv .venv
.venv/Scripts/pip install -e .[dev]
.venv/Scripts/python -m pytest

Contributions follow GUIDELINES.md — benchmarks are mandatory, honesty is policy.