Automatic Guide Generation

AutoContinuous

class AutoContinuous(model, prefix='auto', init_strategy=<function init_to_uniform>)[source]

Bases: numpyro.infer.autoguide.AutoGuide

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

Each derived class implements its own _get_posterior() method.

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

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.
get_base_dist()[source]

Returns the base distribution of the posterior when reparameterized as a TransformedDistribution. This should not depend on the model’s *args, **kwargs.

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. The parameters can be obtained using get_params() method from SVI.
Returns:the transform of posterior distribution
Return type:Transform
get_posterior(params)[source]

Returns the posterior distribution.

Parameters:params (dict) – Current parameters of model and autoguide. The parameters can be obtained using get_params() method from SVI.
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. The parameters can be obtained using get_params() method from SVI.
  • 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. The parameters can be obtained using get_params() method from SVI.
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. The parameters can be obtained using get_params() method from SVI.
  • 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

AutoBNAFNormal

class AutoBNAFNormal(model, prefix='auto', init_strategy=<function init_to_uniform>, num_flows=1, hidden_factors=[8, 8])[source]

Bases: numpyro.infer.autoguide.AutoContinuous

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

Usage:

guide = AutoBNAFNormal(model, num_flows=1, hidden_factors=[50, 50], ...)
svi = SVI(model, guide, ...)

References

  1. Block Neural Autoregressive Flow, Nicola De Cao, Ivan Titov, Wilker Aziz
Parameters:
  • 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.
  • hidden_factors (list) – Hidden layer i has hidden_factors[i] hidden units per input dimension. This corresponds to both \(a\) and \(b\) in reference [1]. The elements of hidden_factors must be integers.
get_base_dist()[source]

AutoDiagonalNormal

class AutoDiagonalNormal(model, prefix='auto', init_strategy=<function init_to_uniform>, init_scale=0.1)[source]

Bases: numpyro.infer.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, ...)
get_base_dist()[source]
get_transform(params)[source]
get_posterior(params)[source]

Returns a diagonal Normal posterior distribution.

median(params)[source]
quantiles(params, quantiles)[source]

AutoMultivariateNormal

class AutoMultivariateNormal(model, prefix='auto', init_strategy=<function init_to_uniform>, init_scale=0.1)[source]

Bases: numpyro.infer.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, ...)
get_base_dist()[source]
get_transform(params)[source]
get_posterior(params)[source]

Returns a multivariate Normal posterior distribution.

median(params)[source]
quantiles(params, quantiles)[source]

AutoIAFNormal

class AutoIAFNormal(model, prefix='auto', init_strategy=<function init_to_uniform>, num_flows=3, hidden_dims=None, skip_connections=False, nonlinearity=(<function elementwise.<locals>.<lambda>>, <function elementwise.<locals>.<lambda>>))[source]

Bases: numpyro.infer.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:
  • 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.
  • hidden_dims (list) – the dimensionality of the hidden units per layer. Defaults to [latent_dim, latent_dim].
  • 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.Elu().
get_base_dist()[source]

AutoLaplaceApproximation

class AutoLaplaceApproximation(model, prefix='auto', init_strategy=<function init_to_uniform>)[source]

Bases: numpyro.infer.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, ...)
get_base_dist()[source]
get_transform(params)[source]
get_posterior(params)[source]

Returns a multivariate Normal posterior distribution.

sample_posterior(rng_key, params, sample_shape=())[source]
median(params)[source]
quantiles(params, quantiles)[source]

AutoLowRankMultivariateNormal

class AutoLowRankMultivariateNormal(model, prefix='auto', init_strategy=<function init_to_uniform>, init_scale=0.1, rank=None)[source]

Bases: numpyro.infer.autoguide.AutoContinuous

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

Usage:

guide = AutoLowRankMultivariateNormal(model, rank=2, ...)
svi = SVI(model, guide, ...)
get_base_dist()[source]
get_transform(params)[source]
get_posterior(params)[source]

Returns a lowrank multivariate Normal posterior distribution.

median(params)[source]
quantiles(params, quantiles)[source]