Automatic Guide Generation¶

AutoGuide¶

class AutoGuide(model, *, prefix='auto', init_loc_fn=<function init_to_uniform>, create_plates=None)[source]¶

Bases: abc.ABC

Base class for automatic guides.

Derived classes must implement the __call__() method.

Parameters

model (callable) – a pyro model
prefix (str) – a prefix that will be prefixed to all param internal sites
init_loc_fn (callable) – A per-site initialization function. See Initialization Strategies section for available functions.
create_plates (callable) – An optional function inputing the same *args,**kwargs as model() and returning a numpyro.plate or iterable of plates. Plates not returned will be created automatically as usual. This is useful for data subsampling.

abstract sample_posterior(rng_key, params, sample_shape=())[source]¶

Generate samples from the approximate posterior over the latent sites in the model.

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) – sample shape of each latent site, 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 value.
Return type: dict

quantiles(params, quantiles)[source]¶

Returns posterior quantiles each latent variable. Example:

print(guide.quantiles(params, [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 an array of quantile values.

Return type

dict

AutoContinuous¶

class AutoContinuous(model, *, prefix='auto', init_loc_fn=<function init_to_uniform>, create_plates=None)[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:

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_loc_fn (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]¶

Generate samples from the approximate posterior over the latent sites in the model.

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) – sample shape of each latent site, defaults to ().

Returns

a dict containing samples drawn the this guide.

Return type

dict

AutoBNAFNormal¶

class AutoBNAFNormal(model, *, prefix='auto', init_loc_fn=<function init_to_uniform>, num_flows=1, hidden_factors=[8, 8], init_strategy=None)[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

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_loc_fn (callable) – A per-site initialization function.
num_flows (int) – the number of flows to be used, defaults to 1.
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]¶: Returns the base distribution of the posterior when reparameterized as a TransformedDistribution. This should not depend on the model’s *args, **kwargs.

AutoDiagonalNormal¶

class AutoDiagonalNormal(model, *, prefix='auto', init_loc_fn=<function init_to_uniform>, init_scale=0.1, init_strategy=None)[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, ...)

scale_constraint = <numpyro.distributions.constraints._SoftplusPositive object>¶

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 a diagonal Normal posterior distribution.

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 value.
Return type: dict

quantiles(params, quantiles)[source]¶

Returns posterior quantiles each latent variable. Example:

print(guide.quantiles(params, [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 an array of quantile values.

Return type

dict

AutoMultivariateNormal¶

class AutoMultivariateNormal(model, *, prefix='auto', init_loc_fn=<function init_to_uniform>, init_scale=0.1, init_strategy=None)[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, ...)

scale_tril_constraint = <numpyro.distributions.constraints._ScaledUnitLowerCholesky object>¶

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 a multivariate Normal posterior distribution.

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 value.
Return type: dict

quantiles(params, quantiles)[source]¶

Returns posterior quantiles each latent variable. Example:

print(guide.quantiles(params, [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 an array of quantile values.

Return type

dict

AutoIAFNormal¶

class AutoIAFNormal(model, *, prefix='auto', init_loc_fn=<function init_to_uniform>, num_flows=3, hidden_dims=None, skip_connections=False, nonlinearity=(<function elementwise.<locals>.<lambda>>, <function elementwise.<locals>.<lambda>>), init_strategy=None)[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_loc_fn (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]¶: Returns the base distribution of the posterior when reparameterized as a TransformedDistribution. This should not depend on the model’s *args, **kwargs.

AutoLaplaceApproximation¶

class AutoLaplaceApproximation(model, *, prefix='auto', init_loc_fn=<function init_to_uniform>, create_plates=None)[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]¶: 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 a multivariate Normal posterior distribution.

sample_posterior(rng_key, params, sample_shape=())[source]¶

Generate samples from the approximate posterior over the latent sites in the model.

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) – sample shape of each latent site, 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 value.
Return type: dict

quantiles(params, quantiles)[source]¶

Returns posterior quantiles each latent variable. Example:

print(guide.quantiles(params, [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 an array of quantile values.

Return type

dict

AutoLowRankMultivariateNormal¶

class AutoLowRankMultivariateNormal(model, *, prefix='auto', init_loc_fn=<function init_to_uniform>, init_scale=0.1, rank=None, init_strategy=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, ...)

scale_constraint = <numpyro.distributions.constraints._SoftplusPositive object>¶

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 a lowrank multivariate Normal posterior distribution.

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 value.
Return type: dict

quantiles(params, quantiles)[source]¶

Returns posterior quantiles each latent variable. Example:

print(guide.quantiles(params, [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 an array of quantile values.

Return type

dict

AutoNormal¶

class AutoNormal(model, *, prefix='auto', init_loc_fn=<function init_to_uniform>, init_scale=0.1, create_plates=None)[source]¶

Bases: numpyro.infer.autoguide.AutoGuide

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

This should be equivalent to AutoDiagonalNormal , but with more convenient site names and with better support for mean field ELBO.

Usage:

guide = AutoNormal(model)
svi = SVI(model, guide, ...)

Parameters

model (callable) – A NumPyro model.
prefix (str) – a prefix that will be prefixed to all param internal sites.
init_loc_fn (callable) – A per-site initialization function. See Initialization Strategies section for available functions.
init_scale (float) – Initial scale for the standard deviation of each (unconstrained transformed) latent variable.
create_plates (callable) – An optional function inputing the same *args,**kwargs as model() and returning a numpyro.plate or iterable of plates. Plates not returned will be created automatically as usual. This is useful for data subsampling.

scale_constraint = <numpyro.distributions.constraints._SoftplusPositive object>¶

sample_posterior(rng_key, params, sample_shape=())[source]¶

Generate samples from the approximate posterior over the latent sites in the model.

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) – sample shape of each latent site, 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 value.
Return type: dict

quantiles(params, quantiles)[source]¶

Returns posterior quantiles each latent variable. Example:

print(guide.quantiles(params, [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 an array of quantile values.

Return type

dict

AutoDelta¶

class AutoDelta(model, *, prefix='auto', init_loc_fn=<function init_to_median>, create_plates=None)[source]¶

Bases: numpyro.infer.autoguide.AutoGuide

This implementation of AutoGuide uses Delta distributions to construct a MAP guide over the entire latent space. The guide does not depend on the model’s *args, **kwargs.

Note

This class does MAP inference in constrained space.

Usage:

guide = AutoDelta(model)
svi = SVI(model, guide, ...)

Parameters

model (callable) – A NumPyro model.
prefix (str) – a prefix that will be prefixed to all param internal sites.
init_loc_fn (callable) – A per-site initialization function. See Initialization Strategies section for available functions.
create_plates (callable) – An optional function inputing the same *args,**kwargs as model() and returning a numpyro.plate or iterable of plates. Plates not returned will be created automatically as usual. This is useful for data subsampling.

sample_posterior(rng_key, params, sample_shape=())[source]¶

Generate samples from the approximate posterior over the latent sites in the model.

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) – sample shape of each latent site, 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 value.
Return type: dict

AutoDAIS¶

class AutoDAIS(model, *, K=4, base_dist='diagonal', eta_init=0.01, eta_max=0.1, gamma_init=0.9, prefix='auto', init_loc_fn=<function init_to_uniform>, init_scale=0.1)[source]¶

Bases: numpyro.infer.autoguide.AutoContinuous

This implementation of AutoDAIS uses Differentiable Annealed Importance Sampling (DAIS) [1, 2] to construct a guide over the entire latent space. Samples from the variational distribution (i.e. guide) are generated using a combination of (uncorrected) Hamiltonian Monte Carlo and Annealed Importance Sampling. The same algorithm is called Uncorrected Hamiltonian Annealing in [1].

Note that AutoDAIS cannot be used in conjuction with data subsampling.

Reference:

MCMC Variational Inference via Uncorrected Hamiltonian Annealing, Tomas Geffner, Justin Domke
Differentiable Annealed Importance Sampling and the Perils of Gradient Noise, Guodong Zhang, Kyle Hsu, Jianing Li, Chelsea Finn, Roger Grosse

Usage:

guide = AutoDAIS(model)
svi = SVI(model, guide, ...)

Parameters

model (callable) – A NumPyro model.
prefix (str) – A prefix that will be prefixed to all param internal sites.
K (int) – A positive integer that controls the number of HMC steps used. Defaults to 4.
base_dist (str) – Controls whether the base Normal variational distribution is parameterized by a “diagonal” covariance matrix or a full-rank covariance matrix parameterized by a lower-triangular “cholesky” factor. Defaults to “diagonal”.
eta_init (float) – The initial value of the step size used in HMC. Defaults to 0.01.
eta_max (float) – The maximum value of the learnable step size used in HMC. Defaults to 0.1.
gamma_init (float) – The initial value of the learnable damping factor used during partial momentum refreshments in HMC. Defaults to 0.9.
init_loc_fn (callable) – A per-site initialization function. See Initialization Strategies section for available functions.
init_scale (float) – Initial scale for the standard deviation of the base variational distribution for each (unconstrained transformed) latent variable. Defaults to 0.1.

sample_posterior(rng_key, params, sample_shape=())[source]¶

Generate samples from the approximate posterior over the latent sites in the model.

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) – sample shape of each latent site, defaults to ().

Returns

a dict containing samples drawn the this guide.

Return type

dict