mixle.inference.mcmc.nuts_numba module

Numba (@njit) No-U-Turn Sampler over an analytic @njit value_and_grad.

A line-for-line port of the numpy NUTS in mixle.inference.mcmc.samplers (recursive tree doubling, U-turn termination, multinomial proposal, dual-averaging step-size adaptation, the _find_reasonable_eps heuristic) compiled with @njit. The (unnormalised) log-target is a PASSED @njit fused value_and_grad(theta) -> (float, ndarray); numba supports first-class jitted functions, so the whole sampler — including the recursive build_tree and its heterogeneous return tuple — runs in nopython mode with no Python-level callback per step.

Contract: the caller supplies an analytic @njit value_and_grad — there is no autodiff here (use the torch or jax backends for that). The win is CPU throughput on analytic-gradient models (no per-step Python dispatch). np.random (seed / standard_normal / exponential / random) is used inside the kernel and is seeded via np.random.seed for reproducibility.

The public entry point nuts_numba() is a thin Python wrapper that loops chains, pools the draws, computes R-hat / ESS, and returns a mixle.inference.mcmc.samplers.MCMCResult — mirroring the numpy facade’s per-chain loop.

nuts_numba(value_and_grad, initial, num_samples=1000, warmup=1000, mass=1.0, target_accept=0.8, max_tree_depth=10, thin=1, seed=None)[source]

No-U-Turn Sampler over an @njit analytic value_and_grad, run in nopython mode.

Parameters:

  • value_and_grad (Callable[[ndarray], tuple[float, ndarray]]) – an @njit-compiled fused callable theta -> (logp, grad) returning the (unnormalised) log target and its analytic gradient. There is no autodiff; supply the gradient (njit-jitted) yourself.

  • initial (ndarray) – starting state, array-like of shape (d,).

  • num_samples (int) – retained draws, adaptation iters, thinning.

  • warmup (int) – retained draws, adaptation iters, thinning.

  • thin (int) – retained draws, adaptation iters, thinning.

  • mass (Any) – diagonal mass matrix (scalar or (d,)).

  • target_accept (float) – NUTS tuning, as in the numpy sampler.

  • max_tree_depth (int) – NUTS tuning, as in the numpy sampler.

  • seed (int | None) – seed for the in-kernel np.random (reproducible per chain).

Returns:

MCMCResult with numpy samples, plus step_size and num_target_evals attributes.

Return type:

MCMCResult