blackjax.vi.multipathfinder

blackjax.vi.multipathfinder#

Classes#

MultipathfinderState

State returned by multi-path Pathfinder.

Functions#

`multi_approximate`(→ tuple[MultipathfinderState, ...)	Multi-path Pathfinder variational inference.
`psis_weights`(→ tuple[blackjax.types.Array, ...)	Compute Pareto-Smoothed Importance Sampling (PSIS) weights.
`as_top_level_api`(→ blackjax.base.VIAlgorithm)	High-level multi-path Pathfinder interface.

Module Contents#

class MultipathfinderState[source]#

State returned by multi-path Pathfinder.

path_states: One PathfinderState per independent L-BFGS run.
samples: Approximate posterior samples drawn from each path’s best approximation, shape (n_paths, num_samples, ...).
logp: Log target density evaluated at the per-path samples, shape (n_paths, num_samples).
logq: Log approximation density at the per-path samples, shape (n_paths, num_samples).

path_states: blackjax.vi.pathfinder.PathfinderState[source]#

samples: blackjax.types.ArrayTree[source]#

logp: blackjax.types.Array[source]#

logq: blackjax.types.Array[source]#

multi_approximate(rng_key: blackjax.types.PRNGKey, logdensity_fn: Callable, initial_positions: blackjax.types.ArrayLikeTree, num_samples: int = 200, *, maxiter: int = 30, maxcor: int = 10, maxls: int = 1000, gtol: float = 1e-08, ftol: float = 1e-05) → tuple[MultipathfinderState, blackjax.vi.pathfinder.PathfinderInfo][source]#

Multi-path Pathfinder variational inference.

Runs single-path Pathfinder independently from each of the supplied initial positions (Algorithm 2 in [ZCGV22]), then collects the per-path samples and log densities needed for importance weighting via psis_weights().

Parameters:

rng_key – PRNG key.
logdensity_fn – (Un-normalised) log density of the target distribution.
initial_positions – Starting points for each L-BFGS run. Must be a pytree where the leading axis indexes the n_paths paths; e.g. an array of shape (n_paths, d).
num_samples – Number of samples drawn per path to estimate ELBO and log weights.
maxiter – Maximum L-BFGS iterations per path.
maxcor – L-BFGS history size.
maxls – Maximum line-search steps per iteration.
gtol – Gradient norm convergence tolerance.
ftol – Function value convergence tolerance.

Returns:

A MultipathfinderState (all path states, per-path samples, and log densities)
and a PathfinderInfo wrapping all per-path PathfinderState objects.

psis_weights(state: MultipathfinderState) → tuple[blackjax.types.Array, blackjax.types.Array][source]#

Compute Pareto-Smoothed Importance Sampling (PSIS) weights.

Thin wrapper around blackjax.util.psis_weights() that extracts the log importance ratios from a MultipathfinderState.

Parameters:

state – Output of multi_approximate().

Returns:

log_weights – Normalised log importance weights, shape (n_paths * num_samples,).
pareto_k – Pareto shape parameter estimate (scalar Array). Values below 0.5 indicate reliable importance sampling; above 0.7 may indicate unreliable estimates.

as_top_level_api(logdensity_fn: Callable) → blackjax.base.VIAlgorithm[source]#

High-level multi-path Pathfinder interface.

Returns a VIAlgorithm whose init runs multi-path Pathfinder and whose sample draws importance-resampled approximate posterior samples using PSIS weights.

Parameters:: logdensity_fn – (Un-normalised) log density of the target distribution.
Return type:: A VIAlgorithm.