Automatic Guide Generation

AutoDiagonalNormal

class AutoDiagonalNormal(model, prefix='auto', init_strategy=functools.partial(<function _init_to_uniform>, radius=2))[source]

Bases: numpyro.contrib.autoguide.AutoContinuous

This implementation of AutoContinuous uses a Normal distribution with a diagonal covariance matrix to construct a guide over the entire latent space. The guide does not depend on the model’s *args, **kwargs.

Usage:

guide = AutoDiagonalNormal(model, ...)
svi = SVI(model, guide, ...)
median(params)[source]
quantiles(params, quantiles)[source]

AutoMultivariateNormal

class AutoMultivariateNormal(model, prefix='auto', init_strategy=functools.partial(<function _init_to_uniform>, radius=2))[source]

Bases: numpyro.contrib.autoguide.AutoContinuous

This implementation of AutoContinuous uses a MultivariateNormal distribution to construct a guide over the entire latent space. The guide does not depend on the model’s *args, **kwargs.

Usage:

guide = AutoMultivariateNormal(model, ...)
svi = SVI(model, guide, ...)
median(params)[source]
quantiles(params, quantiles)[source]

AutoIAFNormal

class AutoIAFNormal(model, prefix='auto', init_strategy=functools.partial(<function _init_to_uniform>, radius=2), num_flows=3, **arn_kwargs)[source]

Bases: numpyro.contrib.autoguide.AutoContinuous

This implementation of AutoContinuous uses a Diagonal Normal distribution transformed via a InverseAutoregressiveTransform to construct a guide over the entire latent space. The guide does not depend on the model’s *args, **kwargs.

Usage:

guide = AutoIAFNormal(model, hidden_dims=[20], skip_connections=True, ...)
svi = SVI(model, guide, ...)
Parameters:
  • rng (jax.random.PRNGKey) – random key to be used as the source of randomness to initialize the guide.
  • model (callable) – a generative model.
  • prefix (str) – a prefix that will be prefixed to all param internal sites.
  • init_strategy (callable) – A per-site initialization function.
  • num_flows (int) – the number of flows to be used, defaults to 3.
  • **arn_kwargs

    keywords for constructing autoregressive neural networks, which includes:

    • hidden_dims (list[int]) - the dimensionality of the hidden units per layer. Defaults to [latent_size, latent_size].
    • skip_connections (bool) - whether to add skip connections from the input to the output of each flow. Defaults to False.
    • nonlinearity (callable) - the nonlinearity to use in the feedforward network. Defaults to jax.experimental.stax.Relu().

AutoLaplaceApproximation

class AutoLaplaceApproximation(model, prefix='auto', init_strategy=functools.partial(<function _init_to_uniform>, radius=2))[source]

Bases: numpyro.contrib.autoguide.AutoContinuous

Laplace approximation (quadratic approximation) approximates the posterior \(\log p(z | x)\) by a multivariate normal distribution in the unconstrained space. Under the hood, it uses Delta distributions to construct a MAP guide over the entire (unconstrained) latent space. Its covariance is given by the inverse of the hessian of \(-\log p(x, z)\) at the MAP point of z.

Usage:

guide = AutoLaplaceApproximation(model, ...)
svi = SVI(model, guide, ...)
sample_posterior(rng, params, sample_shape=())[source]
get_transform(params)[source]
median(params)[source]
quantiles(params, quantiles)[source]

AutoContinuousELBO

class AutoContinuousELBO(num_particles=1)[source]

Bases: numpyro.infer.elbo.ELBO

An ELBO implementation specific to AutoContinuous guides. In those guide, the latent variables of the model are transformed to unconstrained domains. This class provides ELBO of the “transformed” model (i.e. the corresponding model with unconstrained variables) and the guide.

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