Automatic Guide Generation

Warning

The interface for the contrib.autoguide module is experimental, and subject to frequent revisions.

AutoContinuous

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

Bases: numpyro.contrib.autoguide.AutoGuide

Base class for implementations of continuous-valued Automatic Differentiation Variational Inference [1].

Each derived class implements its own _get_transform() method.

Assumes model structure and latent dimension are fixed, and all latent variables are continuous.

Note

We recommend using AutoContinuousELBO as the objective function loss in SVI. In addition, we recommend using sample_posterior() method for drawing posterior samples from the autoguide as it will automatically do any unpacking and transformations required to constrain the samples to the support of the latent sites.

Reference:

  1. Automatic Differentiation Variational Inference, Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, David M. Blei
Parameters:
  • model (callable) – A NumPyro model.
  • prefix (str) – a prefix that will be prefixed to all param internal sites.
  • init_strategy (callable) – A per-site initialization function. See Initialization Strategies section for available functions.
base_dist

Base distribution of the posterior. By default, it is standard normal.

get_transform(params)[source]

Returns the transformation learned by the guide to generate samples from the unconstrained (approximate) posterior.

Parameters:params (dict) – Current parameters of model and autoguide.
Returns:the transform of posterior distribution
Return type:Transform
sample_posterior(rng_key, params, sample_shape=())[source]

Get samples from the learned posterior.

Parameters:
  • rng_key (jax.random.PRNGKey) – random key to be used draw samples.
  • params (dict) – Current parameters of model and autoguide.
  • sample_shape (tuple) – batch shape of each latent sample, defaults to ().
Returns:

a dict containing samples drawn the this guide.

Return type:

dict

median(params)[source]

Returns the posterior median value of each latent variable.

Parameters:params (dict) – A dict containing parameter values.
Returns:A dict mapping sample site name to median tensor.
Return type:dict
quantiles(params, quantiles)[source]

Returns posterior quantiles each latent variable. Example:

print(guide.quantiles(opt_state, [0.05, 0.5, 0.95]))
Parameters:
  • params (dict) – A dict containing parameter values.
  • quantiles (list) – A list of requested quantiles between 0 and 1.
Returns:

A dict mapping sample site name to a list of quantile values.

Return type:

dict

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_key (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_key, 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_key, param_map, model, guide, *args, **kwargs)[source]