mixle.inference.mcmc.nuts_torch module

Torch-native No-U-Turn Sampler.

A device-resident port of the numpy NUTS in mixle.inference.mcmc.samplers: identical algorithm (recursive tree doubling, U-turn termination, multinomial proposal, dual-averaging step-size adaptation), but the leapfrog trajectory and the target evaluation stay in torch tensors on the target’s device. The (unnormalised) log-target is supplied as a torch scalar function logp(theta) -> Tensor[()]; its value_and_grad is built with torch.func and ``torch.compile``d once, so the autograd graph is traced a single time and reused on every leapfrog step instead of being rebuilt (and round-tripped through numpy) per gradient call.

Intended for GPU and large autodiff targets, where staying on-device with a compiled target pays off. On CPU this is typically slower than the numpy sampler (per-op torch dispatch + the tree’s host syncs dominate when the target is cheap), so on CPU prefer the numpy / numba / jax backends. The value here is autodiff without re-tracing the graph every call, plus GPU execution.

nuts_torch(logp, initial, num_samples=1000, warmup=1000, mass=1.0, target_accept=0.8, max_tree_depth=10, thin=1, seed=None, *, compile=True, dtype=None, device=None)[source]

No-U-Turn Sampler over a torch scalar log-target, run entirely on-device.

Parameters:

  • logp (Callable[[Any], Any]) – logp(theta: Tensor[d]) -> Tensor[()] — the (unnormalised) log target.

  • initial (Any) – starting state, array-like or tensor 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 momentum + slice/direction RNG (reproducible).

  • compile (bool) – torch.compile the target (falls back to eager if unavailable).

  • dtype (Any) – torch dtype/device for the trajectory (default float64 / the initial tensor’s device, else CPU).

  • device (Any) – torch dtype/device for the trajectory (default float64 / the initial tensor’s device, else CPU).

Returns:

MCMCResult with samples (numpy), plus step_size and num_target_evals attributes.

Return type:

MCMCResult