blackjax

blackjax#

This module defines a convenient interface to call Blackjax to do inference. You could of course just call pangolin.jax_backend.ancestor_log_prob to get a plain jax function and then call Blackjax yourself. But this module abstracts away all the details.

The module has three layers:

pangolin.blackjax.run_nuts(log_prob, key, initial_state, num_samples)[source]#

Sample from a density using the No-U-Turn Sampler (NUTS).

Warmup uses Blackjax’s window adaptation to tune the step size and mass matrix. Note that num_samples is used both as the length of the adaptation run and as the number of posterior draws.

Parameters:

  • log_prob (Callable) – Log-density of the target distribution, as a jitted jax function over the unconstrained latent space.

  • key (jax.random.PRNGKey) – Randomness source; split internally into warmup and sampling keys.

  • initial_state – Initial position in the unconstrained latent space.

  • num_samples (int) – Number of adaptation steps and number of draws returned.

Returns:

samples – Positions of the returned states, with leading axis num_samples.

Examples

Draw from a 3-dimensional standard normal (no pangolin needed — any jax log-density works):

>>> import jax
>>> import jax.numpy as jnp
>>> log_prob = lambda x: -0.5 * jnp.sum(x ** 2)
>>> key = jax.random.PRNGKey(0)
>>> samps = run_nuts(log_prob, key, jnp.zeros(3), num_samples=10)
>>> samps.shape
(10, 3)
pangolin.blackjax.run_hmc(log_prob, key, initial_state, num_samples, num_integration_steps=60)[source]#

Sample from a density using Hamiltonian Monte Carlo (HMC).

Warmup uses Blackjax’s window adaptation to tune the step size and mass matrix; the trajectory length is fixed via num_integration_steps and is not adapted. As with run_nuts, num_samples is used both as the adaptation length and as the number of posterior draws.

Parameters:

  • log_prob (Callable) – Log-density of the target distribution, as a jitted jax function over the unconstrained latent space.

  • key (jax.random.PRNGKey) – Randomness source; split internally into warmup and sampling keys.

  • initial_state – Initial position in the unconstrained latent space.

  • num_samples (int) – Number of adaptation steps and number of draws returned.

  • num_integration_steps (int, optional) – Fixed number of leapfrog steps per trajectory (default 60). Larger values explore further per iteration at higher cost; if you find yourself needing many steps, prefer NUTS, which adapts trajectory length automatically.

Returns:

samples – Positions of the returned states, with leading axis num_samples.

Examples

Draw from a 3-dimensional standard normal:

>>> import jax
>>> import jax.numpy as jnp
>>> log_prob = lambda x: -0.5 * jnp.sum(x ** 2)
>>> key = jax.random.PRNGKey(0)
>>> samps = run_hmc(log_prob, key, jnp.zeros(3), num_samples=10)
>>> samps.shape
(10, 3)
pangolin.blackjax.run_pathfinder(log_prob, key, initial_state, num_samples, elbo_samples=200, **lbfgs_kwargs)[source]#

Draw approximate posterior samples using Blackjax pathfinder.

Pathfinder runs L-BFGS optimization from the initial position and fits a Gaussian approximation along the optimization path, then draws from it. The draws are i.i.d. samples from the approximation, not a Markov chain: if the true posterior is badly non-Gaussian, results are biased in a way that more draws will not fix. Use run_nuts when exactness matters.

Parameters:

  • log_prob (Callable) – Log-density of the target distribution, as a jitted jax function over the unconstrained latent space.

  • key (jax.random.PRNGKey) – Randomness source; split internally into approximation and sampling keys.

  • initial_state – Initial position in the unconstrained latent space, used as the L-BFGS starting point. Must not be exactly at the mode (zero gradient prevents curvature estimation).

  • num_samples (int) – Number of draws returned from the fitted approximation.

  • elbo_samples (int, optional) – Number of draws used internally to estimate the ELBO along the optimization path (default 200). Controls approximation quality, not the size of the output.

  • **lbfgs_kwargs – Additional keyword arguments forwarded to blackjax.vi.pathfinder.approximate, e.g. maxiter, maxcor, ftol, gtol.

Returns:

samples (jnp.ndarray) – Draws from the Gaussian approximation, with leading axis num_samples.

Examples

Draw approximate samples from a 3-dimensional standard normal (starting away from the mode so L-BFGS can estimate curvature):

>>> import jax
>>> import jax.numpy as jnp
>>> log_prob = lambda x: -0.5 * jnp.sum(x ** 2)
>>> key = jax.random.PRNGKey(0)
>>> samps = run_pathfinder(log_prob, key, jnp.ones(3), num_samples=10)
>>> samps.shape
(10, 3)
pangolin.blackjax.run_meanfield_vi(log_prob, key, initial_state, num_samples, n_iter=500, learning_rate=0.05)[source]#

Approximate a density using mean-field variational inference.

Fits a factorized (diagonal-covariance) Gaussian to the target by stochastic optimization of the ELBO, then draws from it. The approximation cannot capture posterior correlations between latents; for a full-covariance Gaussian use run_fullrank_vi, and for exact sampling use run_nuts.

Parameters:

  • log_prob (Callable) – Log-density of the target distribution, as a jitted jax function over the unconstrained latent space.

  • key (jax.random.PRNGKey) – Randomness source; split internally into optimization and sampling keys.

  • initial_state – Initial position in the unconstrained latent space, used to initialize the variational mean.

  • num_samples (int) – Number of draws returned from the fitted approximation.

  • n_iter (int, optional) – Number of ELBO optimization steps (default 500).

  • learning_rate (float, optional) – Learning rate for the Adam optimizer (default 0.05).

Returns:

samples – Draws from the fitted mean-field Gaussian, with leading axis num_samples.

Examples

Fit a 2-dimensional standard normal and draw 10 samples:

>>> import jax
>>> import jax.numpy as jnp
>>> log_prob = lambda x: -0.5 * jnp.sum(x ** 2)
>>> key = jax.random.PRNGKey(0)
>>> samps = run_meanfield_vi(log_prob, key, jnp.zeros(2), num_samples=10, n_iter=50)
>>> samps.shape
(10, 2)
pangolin.blackjax.run_fullrank_vi(log_prob, key, initial_state, num_samples, n_iter=500, learning_rate=0.05)[source]#

Approximate a density using full-rank variational inference.

Fits a Gaussian with a full (Cholesky-parameterized) covariance to the target by stochastic optimization of the ELBO, then draws from it. Unlike run_meanfield_vi, the approximation captures posterior correlations between latents; it is still a Gaussian approximation, so badly non-Gaussian posteriors are biased in a way that more optimization will not fix.

Parameters:

  • log_prob (Callable) – Log-density of the target distribution, as a jitted jax function over the unconstrained latent space.

  • key (jax.random.PRNGKey) – Randomness source; split internally into optimization and sampling keys.

  • initial_state – Initial position in the unconstrained latent space, used to initialize the variational mean.

  • num_samples (int) – Number of draws returned from the fitted approximation.

  • n_iter (int, optional) – Number of ELBO optimization steps (default 500).

  • learning_rate (float, optional) – Learning rate for the Adam optimizer (default 0.05).

Returns:

samples – Draws from the fitted full-rank Gaussian, with leading axis num_samples.

Examples

Fit a 2-dimensional standard normal and draw 10 samples:

>>> import jax
>>> import jax.numpy as jnp
>>> log_prob = lambda x: -0.5 * jnp.sum(x ** 2)
>>> key = jax.random.PRNGKey(0)
>>> samps = run_fullrank_vi(log_prob, key, jnp.zeros(2), num_samples=10, n_iter=50)
>>> samps.shape
(10, 2)
pangolin.blackjax.run_smc(log_prob, key, initial_state, num_samples, step_size=1.0, inverse_mass_matrix=None, num_integration_steps=10, num_mcmc_steps=10, target_ess=0.5, max_iters=100)[source]#

Sample from a density using adaptive tempered Sequential Monte Carlo.

SMC maintains a population of num_samples particles that is moved from a prior toward the posterior along a temperature schedule chosen adaptively to control the effective sample size. Unlike gradient-based MCMC (NUTS/HMC), SMC can move between separated posterior modes, and unlike VI it is asymptotically exact. It is the most expensive option per sample.

Because tempered SMC is written as prior^tempering x likelihood, a Gaussian pseudo-prior centered at the initial position is used internally, with the target log-density playing the role of the likelihood. initial_state therefore seeds the initial particle cloud.

Parameters:

  • log_prob (Callable) – Log-density of the target distribution, as a jitted jax function over the unconstrained latent space.

  • key (jax.random.PRNGKey) – Randomness source; split internally into initialization and iteration keys.

  • initial_state – Position in the unconstrained latent space around which the initial particle cloud is centered.

  • num_samples (int) – Number of particles, which is also the number of draws returned.

  • step_size (float, optional) – Step size of the inner HMC kernel (default 1.0).

  • inverse_mass_matrix (optional) – Inverse mass matrix of the inner HMC kernel. Defaults to the identity.

  • num_integration_steps (int, optional) – Number of leapfrog steps of the inner HMC kernel (default 10).

  • num_mcmc_steps (int, optional) – Number of MCMC kernel applications per particle per temperature step (default 10).

  • target_ess (float, optional) – Target effective sample size, as a fraction of num_samples, used to choose the next temperature (default 0.5).

  • max_iters (int, optional) – Safety cap on the number of temperature steps (default 100); raises if reached.

Returns:

particles – Final particle population, with leading axis num_samples.

Examples

Sample a 2-dimensional standard normal with 20 particles:

>>> import jax
>>> import jax.numpy as jnp
>>> log_prob = lambda x: -0.5 * jnp.sum(x ** 2)
>>> key = jax.random.PRNGKey(0)
>>> samps = run_smc(log_prob, key, jnp.zeros(2), num_samples=20)
>>> samps.shape
(20, 2)
pangolin.blackjax.run_rwm(log_prob, key, initial_state, num_samples, step_size=1.0)[source]#

Sample from a density using random-walk Metropolis (RWM).

Each step proposes a Gaussian perturbation of the current position and accepts or rejects it via the Metropolis ratio. No gradients are used, so this works on any (even non-differentiable) log-density; the cost is slow, diffusive exploration that scales poorly with dimension. Prefer NUTS whenever gradients are available. There is no adaptation: step_size must be chosen by hand, and the first num_samples steps are discarded as warmup.

Parameters:

  • log_prob (Callable) – Log-density of the target distribution, as a jitted jax function over the unconstrained latent space.

  • key (jax.random.PRNGKey) – Randomness source; split internally into warmup and sampling keys.

  • initial_state – Initial position in the unconstrained latent space.

  • num_samples (int) – Number of draws returned. An equal number of warmup steps is discarded first.

  • step_size (float, optional) – Standard deviation of the Gaussian proposal (default 1.0). Tune so that the acceptance rate is roughly 0.2-0.4.

Returns:

samples – Positions of the returned states, with leading axis num_samples.

Examples

Draw from a 3-dimensional standard normal:

>>> import jax
>>> import jax.numpy as jnp
>>> log_prob = lambda x: -0.5 * jnp.sum(x ** 2)
>>> key = jax.random.PRNGKey(0)
>>> samps = run_rwm(log_prob, key, jnp.zeros(3), num_samples=10)
>>> samps.shape
(10, 3)
pangolin.blackjax.sample_flat(vars, given_vars, given_vals, *, run_inf=<function run_nuts>, bijector_dict={<class 'pangolin.ir.Beta'>: <function <lambda>>, <class 'pangolin.ir.Cauchy'>: None, <class 'pangolin.ir.Dirichlet'>: <function <lambda>>, <class 'pangolin.ir.Exponential'>: <function <lambda>>, <class 'pangolin.ir.Gamma'>: <function <lambda>>, <class 'pangolin.ir.Lognormal'>: <function <lambda>>, <class 'pangolin.ir.MultiNormal'>: None, <class 'pangolin.ir.Normal'>: None, <class 'pangolin.ir.NormalPrec'>: None, <class 'pangolin.ir.StudentT'>: None, <class 'pangolin.ir.Uniform'>: <function <lambda>>, <class 'pangolin.ir.Wishart'>: <function <lambda>>}, deferred=True, **inf_args)[source]#

Run inference over a flat list of variables.

Latent variables are split into:

  • MCMC vars: latents that are ancestors of a given (observed) variable. These are unconstrained via bijectors and sampled by run_inf.

  • Deferred vars: latents with no observed descendants. These are marginalized out of the MCMC target and ancestor-sampled at the end via fill_in, conditioned on posterior draws of their parents.

The split is exact: no MCMC/given var has a deferred parent, the deferred factors integrate to 1 in constrained space, and deferred latents are drawn from their true conditional given the MCMC latents.

pangolin.blackjax.blackjax_calculate(run_inf, frozen=(), **options)[source]#

Wrap an inference driver into a convenient Calculate object.

The driver is always frozen: an engine bound to run_nuts cannot later be called with a different run_inf. All other options are overridable at call time.

Parameters:

  • run_inf – Inference driver, with signature run_inf(log_prob, key, initial_state, num_samples, **options) -> samples, where log_prob is a jitted jax function over the unconstrained latent space, key is a jax PRNGKey, and initial_state is a position in that space. See run_nuts, run_hmc, run_pathfinder for examples. Nothing requires the driver to use Blackjax internally.

  • frozen (Iterable[str]) – Additional option names that cannot be overridden at call time. "run_inf" is always included.

  • **options – Default options forwarded to run_inf (and typically including num_samples).

Returns:

Calculate – Calculator bound to sample_flat with the given driver. Its runtime docstring forwards the driver’s documentation, so help(engine) shows the driver’s parameters.

Return type:

Calculate

Examples

>>> my_nuts = blackjax_calculate(run_nuts, num_samples=500)
>>> my_engine = blackjax_calculate(run_smc, num_particles=500)
pangolin.blackjax.nuts = <pangolin.calculate.Calculate object>#

NUTS inference engine; options are forwarded to run_nuts.

This is a pangolin.calculate.Calculate — see that class for the available methods (sample, E, var, etc.).

pangolin.blackjax.hmc = <pangolin.calculate.Calculate object>#

HMC inference engine; options are forwarded to run_hmc.

This is a pangolin.calculate.Calculate — see that class for the available methods (sample, E, var, etc.).

pangolin.blackjax.pathfinder = <pangolin.calculate.Calculate object>#

Pathfinder inference engine; options are forwarded to run_pathfinder.

This is a pangolin.calculate.Calculate — see that class for the available methods (sample, E, var, etc.).

pangolin.blackjax.meanfield_vi = <pangolin.calculate.Calculate object>#

Mean-field VI inference engine; options are forwarded to run_meanfield_vi.

This is a pangolin.calculate.Calculate — see that class for the available methods (sample, E, var, etc.).

pangolin.blackjax.fullrank_vi = <pangolin.calculate.Calculate object>#

Full-rank VI inference engine; options are forwarded to run_fullrank_vi.

This is a pangolin.calculate.Calculate — see that class for the available methods (sample, E, var, std, sample_arviz).

pangolin.blackjax.smc = <pangolin.calculate.Calculate object>#

Tempered SMC inference engine; options are forwarded to run_smc.

This is a pangolin.calculate.Calculate — see that class for the available methods (sample, E, var, std, sample_arviz).

pangolin.blackjax.rwm = <pangolin.calculate.Calculate object>#

Random-walk Metropolis inference engine; options are forwarded to run_rwm.

This is a pangolin.calculate.Calculate — see that class for the available methods (sample, E, var, std, sample_arviz).

pangolin.blackjax.sample(vars, given_vars=None, given_vals=None, reduce_fn=None, **options)#

Default version of Calculate.sample that draws 1000 samples via NUTS.

Parameters:
pangolin.blackjax.E(vars, given_vars=None, given_vals=None, **options)#

Default version of Calculate.E that uses 1000 samples via NUTS.

Parameters:
pangolin.blackjax.var(vars, given_vars=None, given_vals=None, **options)#

Default version of Calculate.var that uses 1000 samples via NUTS.

Parameters:
pangolin.blackjax.std(vars, given_vars=None, given_vals=None, **options)#

Default version of Calculate.std that uses 1000 samples via NUTS.

Parameters:
pangolin.blackjax.sample_arviz(vars, given_vars=None, given_vals=None, **options)#

Default version of Calculate.sample_arviz that uses 1000 samples via NUTS.

Parameters:

  • vars (dict[str, RV])

  • given_vars (PyTree[pangolin.ir.RV])

  • given_vals (PyTree[ArrayLike])