NUTS

class pmrf.infer.NUTS(num_warmup: int = 1000, target_acceptance_rate: float = 0.8, dense_mass: bool = False, show_progress: bool = True)

Bases: AbstractJointSampler

No-U-Turn Sampler (NUTS) in JAX.

Wrapper around blackjax.nuts.

Automatically handles Stan-style window adaptation for the diagonal or dense inverse mass matrix and step size.

Parameters:

  • num_warmup (int, default=1000) – Number of warmup steps for window adaptation.

  • target_acceptance_rate (float, default=0.8) – Target acceptance rate for step size adaptation.

  • dense_mass (bool, default=False) – If True, estimates a dense mass matrix instead of a diagonal one.

  • show_progress (bool, default=True) – If True, shows progress bars during warmup and sampling loops.

run(logposterior_fn: Callable[[PyTree, Any], Any], y0: PyTree, args: PyTree[Any], key: Array, init_samples: PyTree | None = None, max_steps: int | None = 1000, **kwargs) → SampleResult

Execute the sampling algorithm.

Parameters:

  • logposterior_fn (callable) – A function taking the parameters and args as input and returning the log-posterior probability.

  • y0 (PyTree) – The initial parameters, either for shape reference or as a starting point.

  • args (Any) – Args to pass to fn.

  • key (Array) – A random JAX key.

  • init_samples (PyTree, optional) – An optional batched PyTree the same structure as y0 with initial samples to warm-start the algorithm.

  • max_steps (int, optional) – The maximum number of sampling steps to take. If None, implies there should be no limit.

  • **kwargs – Runtime arguments forward to the solver backend.

Returns:

An instance of pmrf.infer.SampleResult.

Return type:

results

dense_mass: bool = False
num_warmup: int = 1000
show_progress: bool = True
target_acceptance_rate: float = 0.8