VI Inference

VI Inference#

Variational-inference utilities for blayers.

Provides Batched_Trace_ELBO, a drop-in Trace_ELBO replacement that handles minibatching without requiring the model to use numpyro.plate, and svi_run_batched(), an svi.run-style helper that drives it.

Use Batched_Trace_ELBO + svi_run_batched for plate-free batched VI; fall back to standard Trace_ELBO if your model already uses plates.

class blayers.vi_infer.Batched_Trace_ELBO(num_obs, num_particles=1, batch_size=None)[source]#

Bases: ELBO

ELBO estimator for minibatched VI without numpyro.plate.

Behaves like Trace_ELBO but rescales the per-batch log-likelihood by num_obs / batch_size so the gradient is an unbiased estimate of the full-dataset ELBO. Drive it with svi_run_batched().

Assumes all latent variables are global. The whole observed log-likelihood is scaled by num_obs / batch_size and the KL over latents is not rescaled, which is only correct when every latent is shared across observations (the usual case for BLayers: coefficients, scales, embeddings). Models with per-observation (local) latents — e.g. a latent variable sampled once per row — are not supported here; use numpyro.plate with the standard Trace_ELBO instead.

Parameters:

  • num_obs (int) – Total number of observations in the full training set.

  • num_particles (int) – Number of Monte Carlo samples per gradient step.

  • batch_size (int | None) – Minibatch size. If None, inferred from the leading dimension of the first batched kwarg at loss-evaluation time.

Warning

Does not mix with numpyro.plate. A ValueError is raised if a plate is detected in the model trace — the num_obs / batch_size rescaling double-counts plate-subsampled sites, so the ELBO would be silently wrong. Use the standard Trace_ELBO with plates instead.

__init__(num_obs, num_particles=1, batch_size=None)[source]#
Parameters:

  • num_obs (int)

  • num_particles (int)

  • batch_size (int | None)

loss(rng_key, param_map, model, guide, *args, **kwargs)[source]#

Evaluates the ELBO with an estimator that uses num_particles many samples/particles.

Parameters:

  • rng_key (jax.random.key) – random number generator seed.

  • param_map (dict) – dictionary of current parameter values keyed by site name.

  • model (Callable[[...], Any]) – Python callable with NumPyro primitives for the model.

  • guide (Callable[[...], Any]) – Python callable with NumPyro primitives for the guide.

  • args (Any) – arguments to the model / guide (these can possibly vary during the course of fitting).

  • kwargs (Any) – keyword arguments to the model / guide (these can possibly vary during the course of fitting).

Returns:

negative of the Evidence Lower Bound (ELBO) to be minimized.

Return type:

Array

elbo_components(rng_key, param_map, model, guide, *args, **kwargs)[source]#
Parameters:

  • rng_key (Array)

  • param_map (dict[str, Array])

  • model (Callable[[...], Any])

  • guide (Callable[[...], Any])

  • args (Any)

  • kwargs (Any)

Return type:

dict[str, Array]

blayers.vi_infer.svi_run_batched(svi, rng_key, batch_size, num_steps=None, num_epochs=None, **data)[source]#
Parameters:

  • svi (SVI)

  • rng_key (Array)

  • batch_size (int)

  • num_steps (int | None)

  • num_epochs (int | None)

  • data (Array)

Return type:

SVIRunResult