AbstractJointSampler

class pmrf.infer.AbstractJointSampler

Bases: Module

Interface for samplers exploring the joint log-posterior (e.g. MCMC-based NUTS or HMC).

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

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:

A tuple of (pmrf.infer.SampleResult, metrics)`.

Return type:

tuple