Pyro Documentation¶
Installation¶
Getting Started¶
 Install Pyro.
 Learn the basic concepts of Pyro: models and inference.
 Dive in to other tutorials and examples.
Primitives¶

sample
(name, fn, *args, **kwargs)[source]¶ Calls the stochastic function fn with additional sideeffects depending on name and the enclosing context (e.g. an inference algorithm). See Intro I and Intro II for a discussion.
Parameters:  name – name of sample
 fn – distribution class or function
 obs – observed datum (optional; should only be used in context of inference) optionally specified in kwargs
 infer (dict) – Optional dictionary of inference parameters specified in kwargs. See inference documentation for details.
Returns: sample

param
(name, *args, **kwargs)[source]¶ Saves the variable as a parameter in the param store. To interact with the param store or write to disk, see Parameters.
Parameters:  name (str) – name of parameter
 init_tensor (torch.Tensor or callable) – initial tensor or lazy callable that returns a tensor.
For large tensors, it may be cheaper to write e.g.
lambda: torch.randn(100000)
, which will only be evaluated on the initial statement.  constraint (torch.distributions.constraints.Constraint) – torch constraint, defaults to
constraints.real
.  event_dim (int) – (optional) number of rightmost dimensions unrelated to baching. Dimension to the left of this will be considered batch dimensions; if the param statement is inside a subsampled plate, then corresponding batch dimensions of the parameter will be correspondingly subsampled. If unspecified, all dimensions will be considered event dims and no subsampling will be performed.
Returns: parameter
Return type:

module
(name, nn_module, update_module_params=False)[source]¶ Takes a torch.nn.Module and registers its parameters with the ParamStore. In conjunction with the ParamStore save() and load() functionality, this allows the user to save and load modules.
Parameters:  name (str) – name of module
 nn_module (torch.nn.Module) – the module to be registered with Pyro
 update_module_params – determines whether Parameters in the PyTorch module get overridden with the values found in the ParamStore (if any). Defaults to False
Returns: torch.nn.Module

random_module
(name, nn_module, prior, *args, **kwargs)[source]¶ Warning
The random_module primitive is deprecated, and will be removed in a future release. Use
PyroModule
instead to to create Bayesian modules fromtorch.nn.Module
instances. See the Bayesian Regression tutorial for an example.Places a prior over the parameters of the module nn_module. Returns a distribution (callable) over nn.Modules, which upon calling returns a sampled nn.Module.
Parameters:  name (str) – name of pyro module
 nn_module (torch.nn.Module) – the module to be registered with pyro
 prior – pyro distribution, stochastic function, or python dict with parameter names as keys and respective distributions/stochastic functions as values.
Returns: a callable which returns a sampled module

factor
(name, log_factor)[source]¶ Factor statement to add arbitrary log probability factor to a probabilisitic model.
Parameters:  name (str) – Name of the trivial sample
 log_factor (torch.Tensor) – A possibly batched log probability factor.

deterministic
(name, value, event_dim=None)[source]¶ EXPERIMENTAL Deterministic statement to add a
Delta
site with name name and value value to the trace. This is useful when we want to record values which are completely determined by their parents. For example:x = sample("x", dist.Normal(0, 1)) x2 = deterministic("x2", x ** 2)
Note
The site does not affect the model density. This currently converts to a
sample()
statement, but may change in the future.Parameters:  name (str) – Name of the site.
 value (torch.Tensor) – Value of the site.
 event_dim (int) – Optional event dimension, defaults to value.ndim.

class
plate
(name, size=None, subsample_size=None, subsample=None, dim=None, use_cuda=None, device=None)[source]¶ Construct for conditionally independent sequences of variables.
plate
can be used either sequentially as a generator or in parallel as a context manager (formerlyirange
andiarange
, respectively).Sequential
plate
is similar torange()
in that it generates a sequence of values.Vectorized
plate
is similar totorch.arange()
in that it yields an array of indices by which other tensors can be indexed.plate
differs fromtorch.arange()
in that it also informs inference algorithms that the variables being indexed are conditionally independent. To do this,plate
is a provided as context manager rather than a function, and users must guarantee that all computation within anplate
context is conditionally independent:with plate("name", size) as ind: # ...do conditionally independent stuff with ind...
Additionally,
plate
can take advantage of the conditional independence assumptions by subsampling the indices and informing inference algorithms to scale various computed values. This is typically used to subsample minibatches of data:with plate("data", len(data), subsample_size=100) as ind: batch = data[ind] assert len(batch) == 100
By default
subsample_size=False
and this simply yields atorch.arange(0, size)
. If0 < subsample_size <= size
this yields a single random batch of indices of sizesubsample_size
and scales all log likelihood terms bysize/batch_size
, within this context.Warning
This is only correct if all computation is conditionally independent within the context.
Parameters:  name (str) – A unique name to help inference algorithms match
plate
sites between models and guides.  size (int) – Optional size of the collection being subsampled (like stop in builtin range).
 subsample_size (int) – Size of minibatches used in subsampling. Defaults to size.
 subsample (Anything supporting len().) – Optional custom subsample for userdefined subsampling schemes. If specified, then subsample_size will be set to len(subsample).
 dim (int) – An optional dimension to use for this independence index.
If specified,
dim
should be negative, i.e. should index from the right. If not specified,dim
is set to the rightmost dim that is left of all enclosingplate
contexts.  use_cuda (bool) – DEPRECATED, use the device arg instead.
Optional bool specifying whether to use cuda tensors for subsample
and log_prob. Defaults to
torch.Tensor.is_cuda
.  device (str) – Optional keyword specifying which device to place the results of subsample and log_prob on. By default, results are placed on the same device as the default tensor.
Returns: A reusabe context manager yielding a single 1dimensional
torch.Tensor
of indices.Examples:
>>> # This version declares sequential independence and subsamples data: >>> for i in plate('data', 100, subsample_size=10): ... if z[i]: # Control flow in this example prevents vectorization. ... obs = sample('obs_{}'.format(i), dist.Normal(loc, scale), obs=data[i])
>>> # This version declares vectorized independence: >>> with plate('data'): ... obs = sample('obs', dist.Normal(loc, scale), obs=data)
>>> # This version subsamples data in vectorized way: >>> with plate('data', 100, subsample_size=10) as ind: ... obs = sample('obs', dist.Normal(loc, scale), obs=data[ind])
>>> # This wraps a userdefined subsampling method for use in pyro: >>> ind = torch.randint(0, 100, (10,)).long() # custom subsample >>> with plate('data', 100, subsample=ind): ... obs = sample('obs', dist.Normal(loc, scale), obs=data[ind])
>>> # This reuses two different independence contexts. >>> x_axis = plate('outer', 320, dim=1) >>> y_axis = plate('inner', 200, dim=2) >>> with x_axis: ... x_noise = sample("x_noise", dist.Normal(loc, scale)) ... assert x_noise.shape == (320,) >>> with y_axis: ... y_noise = sample("y_noise", dist.Normal(loc, scale)) ... assert y_noise.shape == (200, 1) >>> with x_axis, y_axis: ... xy_noise = sample("xy_noise", dist.Normal(loc, scale)) ... assert xy_noise.shape == (200, 320)
See SVI Part II for an extended discussion.
 name (str) – A unique name to help inference algorithms match

clear_param_store
()[source]¶ Clears the ParamStore. This is especially useful if you’re working in a REPL.

validation_enabled
(is_validate=True)[source]¶ Context manager that is useful when temporarily enabling/disabling validation checks.
Parameters: is_validate (bool) – (optional; defaults to True) temporary validation check override.

enable_validation
(is_validate=True)[source]¶ Enable or disable validation checks in Pyro. Validation checks provide useful warnings and errors, e.g. NaN checks, validating distribution arguments and support values, etc. which is useful for debugging. Since some of these checks may be expensive, we recommend turning this off for mature models.
Parameters: is_validate (bool) – (optional; defaults to True) whether to enable validation checks.

trace
(fn=None, ignore_warnings=False, jit_options=None)[source]¶ Lazy replacement for
torch.jit.trace()
that works with Pyro functions that callpyro.param()
.The actual compilation artifact is stored in the
compiled
attribute of the output. Call diagnostic methods on this attribute.Example:
def model(x): scale = pyro.param("scale", torch.tensor(0.5), constraint=constraints.positive) return pyro.sample("y", dist.Normal(x, scale)) @pyro.ops.jit.trace def model_log_prob_fn(x, y): cond_model = pyro.condition(model, data={"y": y}) tr = pyro.poutine.trace(cond_model).get_trace(x) return tr.log_prob_sum()
Parameters:  fn (callable) – The function to be traced.
 ignore_warnins (bool) – Whether to ignore jit warnings.
 jit_options (dict) – Optional dict of options to pass to
torch.jit.trace()
, e.g.{"optimize": False}
.
Inference¶
In the context of probabilistic modeling, learning is usually called inference. In the particular case of Bayesian inference, this often involves computing (approximate) posterior distributions. In the case of parameterized models, this usually involves some sort of optimization. Pyro supports multiple inference algorithms, with support for stochastic variational inference (SVI) being the most extensive. Look here for more inference algorithms in future versions of Pyro.
See Intro II for a discussion of inference in Pyro.
SVI¶

class
SVI
(model, guide, optim, loss, loss_and_grads=None, num_samples=0, num_steps=0, **kwargs)[source]¶ Bases:
pyro.infer.abstract_infer.TracePosterior
Parameters:  model – the model (callable containing Pyro primitives)
 guide – the guide (callable containing Pyro primitives)
 optim (pyro.optim.PyroOptim) – a wrapper a for a PyTorch optimizer
 loss (pyro.infer.elbo.ELBO) – an instance of a subclass of
ELBO
. Pyro provides three builtin losses:Trace_ELBO
,TraceGraph_ELBO
, andTraceEnum_ELBO
. See theELBO
docs to learn how to implement a custom loss.  num_samples – (DEPRECATED) the number of samples for Monte Carlo posterior approximation
 num_steps – (DEPRECATED) the number of optimization steps to take in
run()
A unified interface for stochastic variational inference in Pyro. The most commonly used loss is
loss=Trace_ELBO()
. See the tutorial SVI Part I for a discussion.
evaluate_loss
(*args, **kwargs)[source]¶ Returns: estimate of the loss Return type: float Evaluate the loss function. Any args or kwargs are passed to the model and guide.

run
(*args, **kwargs)[source]¶ Warning
This method is deprecated, and will be removed in a future release. For inference, use
step()
directly, and for predictions, use thePredictive
class.
ELBO¶

class
ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
object
ELBO
is the toplevel interface for stochastic variational inference via optimization of the evidence lower bound.Most users will not interact with this base class
ELBO
directly; instead they will create instances of derived classes:Trace_ELBO
,TraceGraph_ELBO
, orTraceEnum_ELBO
.Parameters:  num_particles – The number of particles/samples used to form the ELBO (gradient) estimators.
 max_plate_nesting (int) – Optional bound on max number of nested
pyro.plate()
contexts. This is only required when enumerating over sample sites in parallel, e.g. if a site setsinfer={"enumerate": "parallel"}
. If omitted, ELBO may guess a valid value by running the (model,guide) pair once, however this guess may be incorrect if model or guide structure is dynamic.  vectorize_particles (bool) – Whether to vectorize the ELBO computation over num_particles. Defaults to False. This requires static structure in model and guide.
 strict_enumeration_warning (bool) – Whether to warn about possible
misuse of enumeration, i.e. that
pyro.infer.traceenum_elbo.TraceEnum_ELBO
is used iff there are enumerated sample sites.  ignore_jit_warnings (bool) – Flag to ignore warnings from the JIT
tracer. When this is True, all
torch.jit.TracerWarning
will be ignored. Defaults to False.  jit_options (bool) – Optional dict of options to pass to
torch.jit.trace()
, e.g.{"check_trace": True}
.  retain_graph (bool) – Whether to retain autograd graph during an SVI step. Defaults to None (False).
 tail_adaptive_beta (float) – Exponent beta with
1.0 <= beta < 0.0
for use with TraceTailAdaptive_ELBO.
References
[1] Automated Variational Inference in Probabilistic Programming David Wingate, Theo Weber
[2] Black Box Variational Inference, Rajesh Ranganath, Sean Gerrish, David M. Blei

class
Trace_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.elbo.ELBO
A trace implementation of ELBObased SVI. The estimator is constructed along the lines of references [1] and [2]. There are no restrictions on the dependency structure of the model or the guide. The gradient estimator includes partial RaoBlackwellization for reducing the variance of the estimator when nonreparameterizable random variables are present. The RaoBlackwellization is partial in that it only uses conditional independence information that is marked by
plate
contexts. For more finegrained RaoBlackwellization, seeTraceGraph_ELBO
.References
 [1] Automated Variational Inference in Probabilistic Programming,
 David Wingate, Theo Weber
 [2] Black Box Variational Inference,
 Rajesh Ranganath, Sean Gerrish, David M. Blei

loss
(model, guide, *args, **kwargs)[source]¶ Returns: returns an estimate of the ELBO Return type: float Evaluates the ELBO with an estimator that uses num_particles many samples/particles.

class
JitTrace_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.trace_elbo.Trace_ELBO
Like
Trace_ELBO
but usespyro.ops.jit.compile()
to compileloss_and_grads()
.This works only for a limited set of models:
 Models must have static structure.
 Models must not depend on any global data (except the param store).
 All model inputs that are tensors must be passed in via
*args
.  All model inputs that are not tensors must be passed in via
**kwargs
, and compilation will be triggered once per unique**kwargs
.

class
TraceGraph_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.elbo.ELBO
A TraceGraph implementation of ELBObased SVI. The gradient estimator is constructed along the lines of reference [1] specialized to the case of the ELBO. It supports arbitrary dependency structure for the model and guide as well as baselines for nonreparameterizable random variables. Where possible, conditional dependency information as recorded in the
Trace
is used to reduce the variance of the gradient estimator. In particular two kinds of conditional dependency information are used to reduce variance: the sequential order of samples (z is sampled after y => y does not depend on z)
plate
generators
References
 [1] Gradient Estimation Using Stochastic Computation Graphs,
 John Schulman, Nicolas Heess, Theophane Weber, Pieter Abbeel
 [2] Neural Variational Inference and Learning in Belief Networks
 Andriy Mnih, Karol Gregor

loss
(model, guide, *args, **kwargs)[source]¶ Returns: returns an estimate of the ELBO Return type: float Evaluates the ELBO with an estimator that uses num_particles many samples/particles.

loss_and_grads
(model, guide, *args, **kwargs)[source]¶ Returns: returns an estimate of the ELBO Return type: float Computes the ELBO as well as the surrogate ELBO that is used to form the gradient estimator. Performs backward on the latter. Num_particle many samples are used to form the estimators. If baselines are present, a baseline loss is also constructed and differentiated.

class
JitTraceGraph_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.tracegraph_elbo.TraceGraph_ELBO
Like
TraceGraph_ELBO
but usestorch.jit.trace()
to compileloss_and_grads()
.This works only for a limited set of models:
 Models must have static structure.
 Models must not depend on any global data (except the param store).
 All model inputs that are tensors must be passed in via
*args
.  All model inputs that are not tensors must be passed in via
**kwargs
, and compilation will be triggered once per unique**kwargs
.

class
BackwardSampleMessenger
(enum_trace, guide_trace)[source]¶ Bases:
pyro.poutine.messenger.Messenger
Implements forward filtering / backward sampling for sampling from the joint posterior distribution

class
TraceEnum_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.elbo.ELBO
A trace implementation of ELBObased SVI that supports  exhaustive enumeration over discrete sample sites, and  local parallel sampling over any sample site in the guide.
To enumerate over a sample site in the
guide
, mark the site with eitherinfer={'enumerate': 'sequential'}
orinfer={'enumerate': 'parallel'}
. To configure all guide sites at once, useconfig_enumerate()
. To enumerate over a sample site in themodel
, mark the siteinfer={'enumerate': 'parallel'}
and ensure the site does not appear in theguide
.This assumes restricted dependency structure on the model and guide: variables outside of an
plate
can never depend on variables inside thatplate
.
loss
(model, guide, *args, **kwargs)[source]¶ Returns: an estimate of the ELBO Return type: float Estimates the ELBO using
num_particles
many samples (particles).

differentiable_loss
(model, guide, *args, **kwargs)[source]¶ Returns: a differentiable estimate of the ELBO Return type: torch.Tensor Raises: ValueError – if the ELBO is not differentiable (e.g. is identically zero) Estimates a differentiable ELBO using
num_particles
many samples (particles). The result should be infinitely differentiable (as long as underlying derivatives have been implemented).

loss_and_grads
(model, guide, *args, **kwargs)[source]¶ Returns: an estimate of the ELBO Return type: float Estimates the ELBO using
num_particles
many samples (particles). Performs backward on the ELBO of each particle.


class
JitTraceEnum_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.traceenum_elbo.TraceEnum_ELBO
Like
TraceEnum_ELBO
but usespyro.ops.jit.compile()
to compileloss_and_grads()
.This works only for a limited set of models:
 Models must have static structure.
 Models must not depend on any global data (except the param store).
 All model inputs that are tensors must be passed in via
*args
.  All model inputs that are not tensors must be passed in via
**kwargs
, and compilation will be triggered once per unique**kwargs
.

class
TraceMeanField_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.trace_elbo.Trace_ELBO
A trace implementation of ELBObased SVI. This is currently the only ELBO estimator in Pyro that uses analytic KL divergences when those are available.
In contrast to, e.g.,
TraceGraph_ELBO
andTrace_ELBO
this estimator places restrictions on the dependency structure of the model and guide. In particular it assumes that the guide has a meanfield structure, i.e. that it factorizes across the different latent variables present in the guide. It also assumes that all of the latent variables in the guide are reparameterized. This latter condition is satisfied for, e.g., the Normal distribution but is not satisfied for, e.g., the Categorical distribution.Warning
This estimator may give incorrect results if the meanfield condition is not satisfied.
Note for advanced users:
The mean field condition is a sufficient but not necessary condition for this estimator to be correct. The precise condition is that for every latent variable z in the guide, its parents in the model must not include any latent variables that are descendants of z in the guide. Here ‘parents in the model’ and ‘descendants in the guide’ is with respect to the corresponding (statistical) dependency structure. For example, this condition is always satisfied if the model and guide have identical dependency structures.

class
JitTraceMeanField_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.trace_mean_field_elbo.TraceMeanField_ELBO
Like
TraceMeanField_ELBO
but usespyro.ops.jit.trace()
to compileloss_and_grads()
.This works only for a limited set of models:
 Models must have static structure.
 Models must not depend on any global data (except the param store).
 All model inputs that are tensors must be passed in via
*args
.  All model inputs that are not tensors must be passed in via
**kwargs
, and compilation will be triggered once per unique**kwargs
.

class
TraceTailAdaptive_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.trace_elbo.Trace_ELBO
Interface for Stochastic Variational Inference with an adaptive fdivergence as described in ref. [1]. Users should specify num_particles > 1 and vectorize_particles==True. The argument tail_adaptive_beta can be specified to modify how the adaptive fdivergence is constructed. See reference for details.
Note that this interface does not support computing the varational objective itself; rather it only supports computing gradients of the variational objective. Consequently, one might want to use another SVI interface (e.g. RenyiELBO) in order to monitor convergence.
Note that this interface only supports models in which all the latent variables are fully reparameterized. It also does not support data subsampling.
References [1] “Variational Inference with Tailadaptive fDivergence”, Dilin Wang, Hao Liu, Qiang Liu, NeurIPS 2018 https://papers.nips.cc/paper/7816variationalinferencewithtailadaptivefdivergence

class
RenyiELBO
(alpha=0, num_particles=2, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True)[source]¶ Bases:
pyro.infer.elbo.ELBO
An implementation of Renyi’s \(\alpha\)divergence variational inference following reference [1].
In order for the objective to be a strict lower bound, we require \(\alpha \ge 0\). Note, however, that according to reference [1], depending on the dataset \(\alpha < 0\) might give better results. In the special case \(\alpha = 0\), the objective function is that of the important weighted autoencoder derived in reference [2].
Note
Setting \(\alpha < 1\) gives a better bound than the usual ELBO. For \(\alpha = 1\), it is better to use
Trace_ELBO
class because it helps reduce variances of gradient estimations.Parameters:  alpha (float) – The order of \(\alpha\)divergence. Here \(\alpha \neq 1\). Default is 0.
 num_particles – The number of particles/samples used to form the objective (gradient) estimator. Default is 2.
 max_plate_nesting (int) – Bound on max number of nested
pyro.plate()
contexts. Default is infinity.  strict_enumeration_warning (bool) – Whether to warn about possible
misuse of enumeration, i.e. that
TraceEnum_ELBO
is used iff there are enumerated sample sites.
References:
 [1] Renyi Divergence Variational Inference,
 Yingzhen Li, Richard E. Turner
 [2] Importance Weighted Autoencoders,
 Yuri Burda, Roger Grosse, Ruslan Salakhutdinov

class
TraceTMC_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.elbo.ELBO
A tracebased implementation of Tensor Monte Carlo [1] by way of Tensor Variable Elimination [2] that supports:  local parallel sampling over any sample site in the model or guide  exhaustive enumeration over any sample site in the model or guide
To take multiple samples, mark the site with
infer={'enumerate': 'parallel', 'num_samples': N}
. To configure all sites in a model or guide at once, useconfig_enumerate()
. To enumerate or sample a sample site in themodel
, mark the site and ensure the site does not appear in theguide
.This assumes restricted dependency structure on the model and guide: variables outside of an
plate
can never depend on variables inside thatplate
.References
 [1] Tensor Monte Carlo: Particle Methods for the GPU Era,
 Laurence Aitchison (2018)
 [2] Tensor Variable Elimination for Plated Factor Graphs,
 Fritz Obermeyer, Eli Bingham, Martin Jankowiak, Justin Chiu, Neeraj Pradhan, Alexander Rush, Noah Goodman (2019)

differentiable_loss
(model, guide, *args, **kwargs)[source]¶ Returns: a differentiable estimate of the marginal loglikelihood Return type: torch.Tensor Raises: ValueError – if the ELBO is not differentiable (e.g. is identically zero) Computes a differentiable TMC estimate using
num_particles
many samples (particles). The result should be infinitely differentiable (as long as underlying derivatives have been implemented).
Importance¶

class
Importance
(model, guide=None, num_samples=None)[source]¶ Bases:
pyro.infer.abstract_infer.TracePosterior
Parameters:  model – probabilistic model defined as a function
 guide – guide used for sampling defined as a function
 num_samples – number of samples to draw from the guide (default 10)
This method performs posterior inference by importance sampling using the guide as the proposal distribution. If no guide is provided, it defaults to proposing from the model’s prior.

psis_diagnostic
(model, guide, *args, **kwargs)[source]¶ Computes the Pareto tail index k for a model/guide pair using the technique described in [1], which builds on previous work in [2]. If \(0 < k < 0.5\) the guide is a good approximation to the model posterior, in the sense described in [1]. If \(0.5 \le k \le 0.7\), the guide provides a suboptimal approximation to the posterior, but may still be useful in practice. If \(k > 0.7\) the guide program provides a poor approximation to the full posterior, and caution should be used when using the guide. Note, however, that a guide may be a poor fit to the full posterior while still yielding reasonable model predictions. If \(k < 0.0\) the importance weights corresponding to the model and guide appear to be bounded from above; this would be a bizarre outcome for a guide trained via ELBO maximization. Please see [1] for a more complete discussion of how the tail index k should be interpreted.
Please be advised that a large number of samples may be required for an accurate estimate of k.
Note that we assume that the model and guide are both vectorized and have static structure. As is canonical in Pyro, the args and kwargs are passed to the model and guide.
References [1] ‘Yes, but Did It Work?: Evaluating Variational Inference.’ Yuling Yao, Aki Vehtari, Daniel Simpson, Andrew Gelman [2] ‘Pareto Smoothed Importance Sampling.’ Aki Vehtari, Andrew Gelman, Jonah Gabry
Parameters:  model (callable) – the model program.
 guide (callable) – the guide program.
 num_particles (int) – the total number of times we run the model and guide in order to compute the diagnostic. defaults to 1000.
 max_simultaneous_particles – the maximum number of simultaneous samples drawn from the model and guide. defaults to num_particles. num_particles must be divisible by max_simultaneous_particles. compute the diagnostic. defaults to 1000.
 max_plate_nesting (int) – optional bound on max number of nested
pyro.plate()
contexts in the model/guide. defaults to 7.
Returns float: the PSIS diagnostic k

vectorized_importance_weights
(model, guide, *args, **kwargs)[source]¶ Parameters:  model – probabilistic model defined as a function
 guide – guide used for sampling defined as a function
 num_samples – number of samples to draw from the guide (default 1)
 max_plate_nesting (int) – Bound on max number of nested
pyro.plate()
contexts.  normalized (bool) – set to True to return selfnormalized importance weights
Returns: returns a
(num_samples,)
shaped tensor of importance weights and the model and guide traces that produced themVectorized computation of importance weights for models with static structure:
log_weights, model_trace, guide_trace = \ vectorized_importance_weights(model, guide, *args, num_samples=1000, max_plate_nesting=4, normalized=False)
Reweighted WakeSleep¶

class
ReweightedWakeSleep
(num_particles=2, insomnia=1.0, model_has_params=True, num_sleep_particles=None, vectorize_particles=True, max_plate_nesting=inf, strict_enumeration_warning=True)[source]¶ Bases:
pyro.infer.elbo.ELBO
An implementation of Reweighted Wake Sleep following reference [1].
Note
Sampling and log_prob evaluation asymptotic complexity:
 Using waketheta and/or wakephi
 O(num_particles) samples from guide, O(num_particles) log_prob evaluations of model and guide
 Using sleepphi
 O(num_sleep_particles) samples from model, O(num_sleep_particles) log_prob evaluations of guide
 if 1) and 2) are combined,
 O(num_particles) samples from the guide, O(num_sleep_particles) from the model, O(num_particles + num_sleep_particles) log_prob evaluations of the guide, and O(num_particles) evaluations of the model
Note
This is particularly useful for models with stochastic branching, as described in [2].
Note
This returns _two_ losses, one each for (a) the model parameters (theta), computed using the iwae objective, and (b) the guide parameters (phi), computed using (a combination of) the csis objective and a selfnormalized importancesampled version of the csis objective.
Note
In order to enable computing the sleepphi terms, the guide program must have its observations explicitly passed in through the keyworded argument observations. Where the value of the observations is unknown during definition, such as for amortized variational inference, it may be given a default argument as observations=None, and the correct value supplied during learning through svi.step(observations=…).
Warning
Minibatch training is not supported yet.
Parameters:  num_particles (int) – The number of particles/samples used to form the objective (gradient) estimator. Default is 2.
 insomnia – The scaling between the wakephi and sleepphi terms. Default is 1.0 [wakephi]
 model_has_params (bool) – Indicate if model has learnable params. Useful in avoiding extra computation when running in pure sleep mode [csis]. Default is True.
 num_sleep_particles (int) – The number of particles used to form the sleepphi estimator. Matches num_particles by default.
 vectorize_particles (bool) – Whether the traces should be vectorised across num_particles. Default is True.
 max_plate_nesting (int) – Bound on max number of nested
pyro.plate()
contexts. Default is infinity.  strict_enumeration_warning (bool) – Whether to warn about possible
misuse of enumeration, i.e. that
TraceEnum_ELBO
is used iff there are enumerated sample sites.
References:
 [1] Reweighted WakeSleep,
 Jörg Bornschein, Yoshua Bengio
 [2] Revisiting Reweighted WakeSleep for Models with Stochastic Control Flow,
 Tuan Anh Le, Adam R. Kosiorek, N. Siddharth, Yee Whye Teh, Frank Wood
Sequential Monte Carlo¶

class
SMCFilter
(model, guide, num_particles, max_plate_nesting)[source]¶ Bases:
object
SMCFilter
is the toplevel interface for filtering via sequential monte carlo.The model and guide should be objects with two methods:
.init(state, ...)
and.step(state, ...)
, intended to be called first withinit()
, then withstep()
repeatedly. These two methods should have the same signature asSMCFilter
‘sinit()
andstep()
of this class, but with an extra first argumentstate
that should be used to store all tensors that depend on sampled variables. Thestate
will be a dictlike object,SMCState
, with arbitrary keys andtorch.Tensor
values. Models can read and writestate
but guides can only read from it.Inference complexity is
O(len(state) * num_time_steps)
, so to avoid quadratic complexity in Markov models, ensure thatstate
has fixed size.Parameters:  model (object) – probabilistic model with
init
andstep
methods  guide (object) – guide used for sampling, with
init
andstep
methods  num_particles (int) – The number of particles used to form the distribution.
 max_plate_nesting (int) – Bound on max number of nested
pyro.plate()
contexts.

get_empirical
()[source]¶ Returns: a marginal distribution over all state tensors. Return type: a dictionary with keys which are latent variables and values which are Empirical
objects.
 model (object) – probabilistic model with

class
SMCState
(num_particles)[source]¶ Bases:
dict
Dictionarylike object to hold a vectorized collection of tensors to represent all state during inference with
SMCFilter
. During inference, theSMCFilter
resample these tensors.Keys may have arbitrary hashable type. Values must be
torch.Tensor
s.Parameters: num_particles (int) –
Stein Methods¶

class
IMQSteinKernel
(alpha=0.5, beta=0.5, bandwidth_factor=None)[source]¶ Bases:
pyro.infer.svgd.SteinKernel
An IMQ (inverse multiquadratic) kernel for use in the SVGD inference algorithm [1]. The bandwidth of the kernel is chosen from the particles using a simple heuristic as in reference [2]. The kernel takes the form
\(K(x, y) = (\alpha + xy^2/h)^{\beta}\)
where \(\alpha\) and \(\beta\) are userspecified parameters and \(h\) is the bandwidth.
Parameters: Variables: bandwidth_factor (float) – Property that controls the factor by which to scale the bandwidth at each iteration.
References
[1] “Stein Points,” Wilson Ye Chen, Lester Mackey, Jackson Gorham, FrancoisXavier Briol, Chris. J. Oates. [2] “Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm,” Qiang Liu, Dilin Wang

bandwidth_factor
¶


class
RBFSteinKernel
(bandwidth_factor=None)[source]¶ Bases:
pyro.infer.svgd.SteinKernel
A RBF kernel for use in the SVGD inference algorithm. The bandwidth of the kernel is chosen from the particles using a simple heuristic as in reference [1].
Parameters: bandwidth_factor (float) – Optional factor by which to scale the bandwidth, defaults to 1.0. Variables: bandwidth_factor (float) – Property that controls the factor by which to scale the bandwidth at each iteration. References
 [1] “Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm,”
 Qiang Liu, Dilin Wang

bandwidth_factor
¶

class
SVGD
(model, kernel, optim, num_particles, max_plate_nesting, mode='univariate')[source]¶ Bases:
object
A basic implementation of Stein Variational Gradient Descent as described in reference [1].
Parameters:  model – The model (callable containing Pyro primitives). Model must be fully vectorized and may only contain continuous latent variables.
 kernel – a SVGD compatible kernel like
RBFSteinKernel
.  optim (pyro.optim.PyroOptim) – A wrapper for a PyTorch optimizer.
 num_particles (int) – The number of particles used in SVGD.
 max_plate_nesting (int) – The max number of nested
pyro.plate()
contexts in the model.  mode (str) – Whether to use a Kernelized Stein Discrepancy that makes use of multivariate test functions (as in [1]) or univariate test functions (as in [2]). Defaults to univariate.
Example usage:
from pyro.infer import SVGD, RBFSteinKernel from pyro.optim import Adam kernel = RBFSteinKernel() adam = Adam({"lr": 0.1}) svgd = SVGD(model, kernel, adam, num_particles=50, max_plate_nesting=0) for step in range(500): svgd.step(model_arg1, model_arg2) final_particles = svgd.get_named_particles()
References
 [1] “Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm,”
 Qiang Liu, Dilin Wang
 [2] “Kernelized Complete Conditional Stein Discrepancy,”
 Raghav Singhal, Saad Lahlou, Rajesh Ranganath

class
SteinKernel
[source]¶ Bases:
object
Abstract class for kernels used in the
SVGD
inference algorithm.
log_kernel_and_grad
(particles)[source]¶ Compute the component kernels and their gradients.
Parameters: particles – a tensor with shape (N, D) Returns: A pair (log_kernel, kernel_grad) where log_kernel is a (N, N, D)shaped tensor equal to the logarithm of the kernel and kernel_grad is a (N, N, D)shaped tensor where the entry (n, m, d) represents the derivative of log_kernel w.r.t. x_{m,d}, where x_{m,d} is the d^th dimension of particle m.

Likelihood free methods¶

class
EnergyDistance
(beta=1.0, prior_scale=0.0, num_particles=2, max_plate_nesting=inf)[source]¶ Bases:
object
Posterior predictive energy distance [1,2] with optional Bayesian regularization by the prior.
Let p(x,z)=p(z) p(xz) be the model, q(zx) be the guide. Then given data x and drawing an iid pair of samples \((Z,X)\) and \((Z',X')\) (where Z is latent and X is the posterior predictive),
\[\begin{split}& Z \sim q(zx); \quad X \sim p(xZ) \\ & Z' \sim q(zx); \quad X' \sim p(xZ') \\ & loss = \mathbb E_X \Xx\^\beta  \frac 1 2 \mathbb E_{X,X'}\XX'\^\beta  \lambda \mathbb E_Z \log p(Z)\end{split}\]This is a likelihoodfree inference algorithm, and can be used for likelihoods without tractable density functions. The \(\beta\) energy distance is a robust loss functions, and is well defined for any distribution with finite fractional moment \(\mathbb E[\X\^\beta]\).
This requires static model structure, a fully reparametrized guide, and reparametrized likelihood distributions in the model. Model latent distributions may be nonreparametrized.
References
 [1] Gabor J. Szekely, Maria L. Rizzo (2003)
 Energy Statistics: A Class of Statistics Based on Distances.
 [2] Tilmann Gneiting, Adrian E. Raftery (2007)
 Strictly Proper Scoring Rules, Prediction, and Estimation. https://www.stat.washington.edu/raftery/Research/PDF/Gneiting2007jasa.pdf
Parameters:  beta (float) – Exponent \(\beta\) from [1,2]. The loss function is
strictly proper for distributions with finite \(beta\)absolute moment
\(E[\X\^\beta]\). Thus for heavy tailed distributions
beta
should be small, e.g. forCauchy
distributions, \(\beta<1\) is strictly proper. Defaults to 1. Must be in the open interval (0,2).  prior_scale (float) – Nonnegative scale for prior regularization. Model parameters are trained only if this is positive. If zero (default), then model log densities will not be computed (guide log densities are never computed).
 num_particles (int) – The number of particles/samples used to form the gradient estimators. Must be at least 2.
 max_plate_nesting (int) – Optional bound on max number of nested
pyro.plate()
contexts. If omitted, this will guess a valid value by running the (model,guide) pair once.
Discrete Inference¶

infer_discrete
(fn=None, first_available_dim=None, temperature=1)[source]¶ A poutine that samples discrete sites marked with
site["infer"]["enumerate"] = "parallel"
from the posterior, conditioned on observations.Example:
@infer_discrete(first_available_dim=1, temperature=0) @config_enumerate def viterbi_decoder(data, hidden_dim=10): transition = 0.3 / hidden_dim + 0.7 * torch.eye(hidden_dim) means = torch.arange(float(hidden_dim)) states = [0] for t in pyro.markov(range(len(data))): states.append(pyro.sample("states_{}".format(t), dist.Categorical(transition[states[1]]))) pyro.sample("obs_{}".format(t), dist.Normal(means[states[1]], 1.), obs=data[t]) return states # returns maximum likelihood states
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 first_available_dim (int) – The first tensor dimension (counting from the right) that is available for parallel enumeration. This dimension and all dimensions left may be used internally by Pyro. This should be a negative integer.
 temperature (int) – Either 1 (sample via forwardfilter backwardsample) or 0 (optimize via Viterbilike MAP inference). Defaults to 1 (sample).

class
TraceEnumSample_ELBO
(num_particles=1, max_plate_nesting=inf, max_iarange_nesting=None, vectorize_particles=False, strict_enumeration_warning=True, ignore_jit_warnings=False, jit_options=None, retain_graph=None, tail_adaptive_beta=1.0)[source]¶ Bases:
pyro.infer.traceenum_elbo.TraceEnum_ELBO
This extends
TraceEnum_ELBO
to make it cheaper to sample from discrete latent states during SVI.The following are equivalent but the first is cheaper, sharing work between the computations of
loss
andz
:# Version 1. elbo = TraceEnumSample_ELBO(max_plate_nesting=1) loss = elbo.loss(*args, **kwargs) z = elbo.sample_saved() # Version 2. elbo = TraceEnum_ELBO(max_plate_nesting=1) loss = elbo.loss(*args, **kwargs) guide_trace = poutine.trace(guide).get_trace(*args, **kwargs) z = infer_discrete(poutine.replay(model, guide_trace), first_available_dim=2)(*args, **kwargs)
Inference Utilities¶

class
Predictive
(model, posterior_samples=None, guide=None, num_samples=None, return_sites=(), parallel=False)[source]¶ Bases:
torch.nn.modules.module.Module
EXPERIMENTAL class used to construct predictive distribution. The predictive distribution is obtained by running the model conditioned on latent samples from posterior_samples. If a guide is provided, then posterior samples from all the latent sites are also returned.
Warning
The interface for the
Predictive
class is experimental, and might change in the future.Parameters:  model – Python callable containing Pyro primitives.
 posterior_samples (dict) – dictionary of samples from the posterior.
 guide (callable) – optional guide to get posterior samples of sites not present in posterior_samples.
 num_samples (int) – number of samples to draw from the predictive distribution.
This argument has no effect if
posterior_samples
is nonempty, in which case, the leading dimension size of samples inposterior_samples
is used.  return_sites (list, tuple, or set) – sites to return; by default only sample sites not present in posterior_samples are returned.
 parallel (bool) – predict in parallel by wrapping the existing model
in an outermost plate messenger. Note that this requires that the model has
all batch dims correctly annotated via
plate
. Default is False.

call
(*args, **kwargs)[source]¶ Method that calls
forward()
and returns parameter values of the guide as a tuple instead of a dict, which is a requirement for JIT tracing. Unlikeforward()
, this method can be traced bytorch.jit.trace_module()
.Warning
This method may be removed once PyTorch JIT tracer starts accepting dict as valid return types. See issue.

forward
(*args, **kwargs)[source]¶ Returns dict of samples from the predictive distribution. By default, only sample sites not contained in posterior_samples are returned. This can be modified by changing the return_sites keyword argument of this
Predictive
instance.Parameters:  args – model arguments.
 kwargs – model keyword arguments.

class
EmpiricalMarginal
(trace_posterior, sites=None, validate_args=None)[source]¶ Bases:
pyro.distributions.empirical.Empirical
Marginal distribution over a single site (or multiple, provided they have the same shape) from the
TracePosterior
’s model.Note
If multiple sites are specified, they must have the same tensor shape. Samples from each site will be stacked and stored within a single tensor. See
Empirical
. To hold the marginal distribution of sites having different shapes, useMarginals
instead.Parameters:  trace_posterior (TracePosterior) – a
TracePosterior
instance representing a Monte Carlo posterior.  sites (list) – optional list of sites for which we need to generate the marginal distribution.
 trace_posterior (TracePosterior) – a

class
Marginals
(trace_posterior, sites=None, validate_args=None)[source]¶ Bases:
object
Holds the marginal distribution over one or more sites from the
TracePosterior
’s model. This is a convenience container class, which can be extended byTracePosterior
subclasses. e.g. for implementing diagnostics.Parameters:  trace_posterior (TracePosterior) – a TracePosterior instance representing a Monte Carlo posterior.
 sites (list) – optional list of sites for which we need to generate the marginal distribution.

empirical
¶ A dictionary of sites’ names and their corresponding
EmpiricalMarginal
distribution.Type: OrderedDict

support
(flatten=False)[source]¶ Gets support of this marginal distribution.
Parameters: flatten (bool) – A flag to decide if we want to flatten batch_shape when the marginal distribution is collected from the posterior with num_chains > 1
. Defaults to False.Returns: a dict with keys are sites’ names and values are sites’ supports. Return type: OrderedDict

class
TracePosterior
(num_chains=1)[source]¶ Bases:
object
Abstract TracePosterior object from which posterior inference algorithms inherit. When run, collects a bag of execution traces from the approximate posterior. This is designed to be used by other utility classes like EmpiricalMarginal, that need access to the collected execution traces.

information_criterion
(pointwise=False)[source]¶ Computes information criterion of the model. Currently, returns only “Widely Applicable/WatanabeAkaike Information Criterion” (WAIC) and the corresponding effective number of parameters.
Reference:
[1] Practical Bayesian model evaluation using leaveoneout crossvalidation and WAIC, Aki Vehtari, Andrew Gelman, and Jonah Gabry
Parameters: pointwise (bool) – a flag to decide if we want to get a vectorized WAIC or not. When pointwise=False
, returns the sum.Returns: a dictionary containing values of WAIC and its effective number of parameters. Return type: OrderedDict


class
TracePredictive
(model, posterior, num_samples, keep_sites=None)[source]¶ Bases:
pyro.infer.abstract_infer.TracePosterior
Warning
This class is deprecated and will be removed in a future release. Use the
Predictive
class instead.Generates and holds traces from the posterior predictive distribution, given model execution traces from the approximate posterior. This is achieved by constraining latent sites to randomly sampled parameter values from the model execution traces and running the model forward to generate traces with new response (“_RETURN”) sites. :param model: arbitrary Python callable containing Pyro primitives. :param TracePosterior posterior: trace posterior instance holding samples from the model’s approximate posterior. :param int num_samples: number of samples to generate. :param keep_sites: The sites which should be sampled from posterior distribution (default: all)
MCMC¶
MCMC¶

class
MCMC
(kernel, num_samples, warmup_steps=None, initial_params=None, num_chains=1, hook_fn=None, mp_context=None, disable_progbar=False, disable_validation=True, transforms=None)[source]¶ Bases:
object
Wrapper class for Markov Chain Monte Carlo algorithms. Specific MCMC algorithms are TraceKernel instances and need to be supplied as a
kernel
argument to the constructor.Note
The case of num_chains > 1 uses python multiprocessing to run parallel chains in multiple processes. This goes with the usual caveats around multiprocessing in python, e.g. the model used to initialize the
kernel
must be serializable via pickle, and the performance / constraints will be platform dependent (e.g. only the “spawn” context is available in Windows). This has also not been extensively tested on the Windows platform.Parameters:  kernel – An instance of the
TraceKernel
class, which when given an execution trace returns another sample trace from the target (posterior) distribution.  num_samples (int) – The number of samples that need to be generated, excluding the samples discarded during the warmup phase.
 warmup_steps (int) – Number of warmup iterations. The samples generated during the warmup phase are discarded. If not provided, default is half of num_samples.
 num_chains (int) – Number of MCMC chains to run in parallel. Depending on whether num_chains is 1 or more than 1, this class internally dispatches to either _UnarySampler or _MultiSampler.
 initial_params (dict) – dict containing initial tensors in unconstrained space to initiate the markov chain. The leading dimension’s size must match that of num_chains. If not specified, parameter values will be sampled from the prior.
 hook_fn – Python callable that takes in (kernel, samples, stage, i) as arguments. stage is either sample or warmup and i refers to the i’th sample for the given stage. This can be used to implement additional logging, or more generally, run arbitrary code per generated sample.
 mp_context (str) – Multiprocessing context to use when num_chains > 1. Only applicable for Python 3.5 and above. Use mp_context=”spawn” for CUDA.
 disable_progbar (bool) – Disable progress bar and diagnostics update.
 disable_validation (bool) – Disables distribution validation check. This is disabled by default, since divergent transitions will lead to exceptions. Switch to True for debugging purposes.
 transforms (dict) – dictionary that specifies a transform for a sample site with constrained support to unconstrained space.

diagnostics
()[source]¶ Gets some diagnostics statistics such as effective sample size, split GelmanRubin, or divergent transitions from the sampler.

get_samples
(num_samples=None, group_by_chain=False)[source]¶ Get samples from the MCMC run, potentially resampling with replacement.
Parameters: Returns: dictionary of samples keyed by site name.

run
¶

summary
(prob=0.9)[source]¶ Prints a summary table displaying diagnostics of samples obtained from posterior. The diagnostics displayed are mean, standard deviation, median, the 90% Credibility Interval,
effective_sample_size()
,split_gelman_rubin()
.Parameters: prob (float) – the probability mass of samples within the credibility interval.
 kernel – An instance of the
HMC¶

class
HMC
(model=None, potential_fn=None, step_size=1, trajectory_length=None, num_steps=None, adapt_step_size=True, adapt_mass_matrix=True, full_mass=False, transforms=None, max_plate_nesting=None, jit_compile=False, jit_options=None, ignore_jit_warnings=False, target_accept_prob=0.8)[source]¶ Bases:
pyro.infer.mcmc.mcmc_kernel.MCMCKernel
Simple Hamiltonian Monte Carlo kernel, where
step_size
andnum_steps
need to be explicitly specified by the user.References
[1] MCMC Using Hamiltonian Dynamics, Radford M. Neal
Parameters:  model – Python callable containing Pyro primitives.
 potential_fn – Python callable calculating potential energy with input is a dict of real support parameters.
 step_size (float) – Determines the size of a single step taken by the verlet integrator while computing the trajectory using Hamiltonian dynamics. If not specified, it will be set to 1.
 trajectory_length (float) – Length of a MCMC trajectory. If not
specified, it will be set to
step_size x num_steps
. In casenum_steps
is not specified, it will be set to \(2\pi\).  num_steps (int) – The number of discrete steps over which to simulate
Hamiltonian dynamics. The state at the end of the trajectory is
returned as the proposal. This value is always equal to
int(trajectory_length / step_size)
.  adapt_step_size (bool) – A flag to decide if we want to adapt step_size during warmup phase using Dual Averaging scheme.
 adapt_mass_matrix (bool) – A flag to decide if we want to adapt mass matrix during warmup phase using Welford scheme.
 full_mass (bool) – A flag to decide if mass matrix is dense or diagonal.
 transforms (dict) – Optional dictionary that specifies a transform
for a sample site with constrained support to unconstrained space. The
transform should be invertible, and implement log_abs_det_jacobian.
If not specified and the model has sites with constrained support,
automatic transformations will be applied, as specified in
torch.distributions.constraint_registry
.  max_plate_nesting (int) – Optional bound on max number of nested
pyro.plate()
contexts. This is required if model contains discrete sample sites that can be enumerated over in parallel.  jit_compile (bool) – Optional parameter denoting whether to use the PyTorch JIT to trace the log density computation, and use this optimized executable trace in the integrator.
 jit_options (dict) – A dictionary contains optional arguments for
torch.jit.trace()
function.  ignore_jit_warnings (bool) – Flag to ignore warnings from the JIT
tracer when
jit_compile=True
. Default is False.  target_accept_prob (float) – Increasing this value will lead to a smaller step size, hence the sampling will be slower and more robust. Default to 0.8.
Note
Internally, the mass matrix will be ordered according to the order of the names of latent variables, not the order of their appearance in the model.
Example:
>>> true_coefs = torch.tensor([1., 2., 3.]) >>> data = torch.randn(2000, 3) >>> dim = 3 >>> labels = dist.Bernoulli(logits=(true_coefs * data).sum(1)).sample() >>> >>> def model(data): ... coefs_mean = torch.zeros(dim) ... coefs = pyro.sample('beta', dist.Normal(coefs_mean, torch.ones(3))) ... y = pyro.sample('y', dist.Bernoulli(logits=(coefs * data).sum(1)), obs=labels) ... return y >>> >>> hmc_kernel = HMC(model, step_size=0.0855, num_steps=4) >>> mcmc = MCMC(hmc_kernel, num_samples=500, warmup_steps=100) >>> mcmc.run(data) >>> mcmc.get_samples()['beta'].mean(0) # doctest: +SKIP tensor([ 0.9819, 1.9258, 2.9737])

initial_params
¶

inverse_mass_matrix
¶

num_steps
¶

step_size
¶
NUTS¶

class
NUTS
(model=None, potential_fn=None, step_size=1, adapt_step_size=True, adapt_mass_matrix=True, full_mass=False, use_multinomial_sampling=True, transforms=None, max_plate_nesting=None, jit_compile=False, jit_options=None, ignore_jit_warnings=False, target_accept_prob=0.8, max_tree_depth=10)[source]¶ Bases:
pyro.infer.mcmc.hmc.HMC
NoUTurn Sampler kernel, which provides an efficient and convenient way to run Hamiltonian Monte Carlo. The number of steps taken by the integrator is dynamically adjusted on each call to
sample
to ensure an optimal length for the Hamiltonian trajectory [1]. As such, the samples generated will typically have lower autocorrelation than those generated by theHMC
kernel. Optionally, the NUTS kernel also provides the ability to adapt step size during the warmup phase.Refer to the baseball example to see how to do Bayesian inference in Pyro using NUTS.
References
 [1] The NoUturn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo,
 Matthew D. Hoffman, and Andrew Gelman.
 [2] A Conceptual Introduction to Hamiltonian Monte Carlo,
 Michael Betancourt
 [3] Slice Sampling,
 Radford M. Neal
Parameters:  model – Python callable containing Pyro primitives.
 potential_fn – Python callable calculating potential energy with input is a dict of real support parameters.
 step_size (float) – Determines the size of a single step taken by the verlet integrator while computing the trajectory using Hamiltonian dynamics. If not specified, it will be set to 1.
 adapt_step_size (bool) – A flag to decide if we want to adapt step_size during warmup phase using Dual Averaging scheme.
 adapt_mass_matrix (bool) – A flag to decide if we want to adapt mass matrix during warmup phase using Welford scheme.
 full_mass (bool) – A flag to decide if mass matrix is dense or diagonal.
 use_multinomial_sampling (bool) – A flag to decide if we want to sample candidates along its trajectory using “multinomial sampling” or using “slice sampling”. Slice sampling is used in the original NUTS paper [1], while multinomial sampling is suggested in [2]. By default, this flag is set to True. If it is set to False, NUTS uses slice sampling.
 transforms (dict) – Optional dictionary that specifies a transform
for a sample site with constrained support to unconstrained space. The
transform should be invertible, and implement log_abs_det_jacobian.
If not specified and the model has sites with constrained support,
automatic transformations will be applied, as specified in
torch.distributions.constraint_registry
.  max_plate_nesting (int) – Optional bound on max number of nested
pyro.plate()
contexts. This is required if model contains discrete sample sites that can be enumerated over in parallel.  jit_compile (bool) – Optional parameter denoting whether to use the PyTorch JIT to trace the log density computation, and use this optimized executable trace in the integrator.
 jit_options (dict) – A dictionary contains optional arguments for
torch.jit.trace()
function.  ignore_jit_warnings (bool) – Flag to ignore warnings from the JIT
tracer when
jit_compile=True
. Default is False.  target_accept_prob (float) – Target acceptance probability of step size adaptation scheme. Increasing this value will lead to a smaller step size, so the sampling will be slower but more robust. Default to 0.8.
 max_tree_depth (int) – Max depth of the binary tree created during the doubling scheme of NUTS sampler. Default to 10.
Example:
>>> true_coefs = torch.tensor([1., 2., 3.]) >>> data = torch.randn(2000, 3) >>> dim = 3 >>> labels = dist.Bernoulli(logits=(true_coefs * data).sum(1)).sample() >>> >>> def model(data): ... coefs_mean = torch.zeros(dim) ... coefs = pyro.sample('beta', dist.Normal(coefs_mean, torch.ones(3))) ... y = pyro.sample('y', dist.Bernoulli(logits=(coefs * data).sum(1)), obs=labels) ... return y >>> >>> nuts_kernel = NUTS(model, adapt_step_size=True) >>> mcmc = MCMC(nuts_kernel, num_samples=500, warmup_steps=300) >>> mcmc.run(data) >>> mcmc.get_samples()['beta'].mean(0) # doctest: +SKIP tensor([ 0.9221, 1.9464, 2.9228])
Utilities¶

initialize_model
(model, model_args=(), model_kwargs={}, transforms=None, max_plate_nesting=None, jit_compile=False, jit_options=None, skip_jit_warnings=False, num_chains=1)[source]¶ Given a Python callable with Pyro primitives, generates the following modelspecific properties needed for inference using HMC/NUTS kernels:
 initial parameters to be sampled using a HMC kernel,
 a potential function whose input is a dict of parameters in unconstrained space,
 transforms to transform latent sites of model to unconstrained space,
 a prototype trace to be used in MCMC to consume traces from sampled parameters.
Parameters:  model – a Pyro model which contains Pyro primitives.
 model_args (tuple) – optional args taken by model.
 model_kwargs (dict) – optional kwargs taken by model.
 transforms (dict) – Optional dictionary that specifies a transform
for a sample site with constrained support to unconstrained space. The
transform should be invertible, and implement log_abs_det_jacobian.
If not specified and the model has sites with constrained support,
automatic transformations will be applied, as specified in
torch.distributions.constraint_registry
.  max_plate_nesting (int) – Optional bound on max number of nested
pyro.plate()
contexts. This is required if model contains discrete sample sites that can be enumerated over in parallel.  jit_compile (bool) – Optional parameter denoting whether to use the PyTorch JIT to trace the log density computation, and use this optimized executable trace in the integrator.
 jit_options (dict) – A dictionary contains optional arguments for
torch.jit.trace()
function.  ignore_jit_warnings (bool) – Flag to ignore warnings from the JIT
tracer when
jit_compile=True
. Default is False.  num_chains (int) – Number of parallel chains. If num_chains > 1, the returned initial_params will be a list with num_chains elements.
Returns: a tuple of (initial_params, potential_fn, transforms, prototype_trace)
Automatic Guide Generation¶
AutoGuide¶

class
AutoGuide
(model)[source]¶ Bases:
pyro.nn.module.PyroModule
Base class for automatic guides.
Derived classes must implement the
forward()
method, with the same*args, **kwargs
as the basemodel
.Auto guides can be used individually or combined in an
AutoGuideList
object.Parameters: model (callable) – a pyro model 
call
(*args, **kwargs)[source]¶ Method that calls
forward()
and returns parameter values of the guide as a tuple instead of a dict, which is a requirement for JIT tracing. Unlikeforward()
, this method can be traced bytorch.jit.trace_module()
.Warning
This method may be removed once PyTorch JIT tracer starts accepting dict as valid return types. See issue <https://github.com/pytorch/pytorch/issues/27743>_.

median
(*args, **kwargs)[source]¶ Returns the posterior median value of each latent variable.
Returns: A dict mapping sample site name to median tensor. Return type: dict

model
¶

AutoGuideList¶

class
AutoGuideList
(model)[source]¶ Bases:
pyro.infer.autoguide.guides.AutoGuide
,torch.nn.modules.container.ModuleList
Container class to combine multiple automatic guides.
Example usage:
guide = AutoGuideList(my_model) guide.add(AutoDiagonalNormal(poutine.block(model, hide=["assignment"]))) guide.add(AutoDiscreteParallel(poutine.block(model, expose=["assignment"]))) svi = SVI(model, guide, optim, Trace_ELBO())
Parameters: model (callable) – a Pyro model 
append
(part)[source]¶ Add an automatic guide for part of the model. The guide should have been created by blocking the model to restrict to a subset of sample sites. No two parts should operate on any one sample site.
Parameters: part (AutoGuide or callable) – a partial guide to add

AutoCallable¶

class
AutoCallable
(model, guide, median=<function AutoCallable.<lambda>>)[source]¶ Bases:
pyro.infer.autoguide.guides.AutoGuide
AutoGuide
wrapper for simple callable guides.This is used internally for composing autoguides with custom userdefined guides that are simple callables, e.g.:
def my_local_guide(*args, **kwargs): ... guide = AutoGuideList(model) guide.add(AutoDelta(poutine.block(model, expose=['my_global_param'])) guide.add(my_local_guide) # automatically wrapped in an AutoCallable
To specify a median callable, you can instead:
def my_local_median(*args, **kwargs) ... guide.add(AutoCallable(model, my_local_guide, my_local_median))
For more complex guides that need e.g. access to plates, users should instead subclass
AutoGuide
.Parameters:  model (callable) – a Pyro model
 guide (callable) – a Pyro guide (typically over only part of the model)
 median (callable) – an optional callable returning a dict mapping sample site name to computed median tensor.
AutoDelta¶

class
AutoDelta
(model, init_loc_fn=<function init_to_median>)[source]¶ Bases:
pyro.infer.autoguide.guides.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, ...)
Latent variables are initialized using
init_loc_fn()
. To change the default behavior, create a custominit_loc_fn()
as described in Initialization , for example:def my_init_fn(site): if site["name"] == "level": return torch.tensor([1., 0., 1.]) if site["name"] == "concentration": return torch.ones(k) return init_to_sample(site)
Parameters:  model (callable) – A Pyro model.
 init_loc_fn (callable) – A persite initialization function. See Initialization section for available functions.
AutoContinuous¶

class
AutoContinuous
(model, init_loc_fn=<function init_to_median>)[source]¶ Bases:
pyro.infer.autoguide.guides.AutoGuide
Base class for implementations of continuousvalued Automatic Differentiation Variational Inference [1].
This uses
torch.distributions.transforms
to transform each constrained latent variable to an unconstrained space, then concatenate all variables into a single unconstrained latent variable. Each derived class implements aget_posterior()
method returning a distribution over this single unconstrained latent variable.Assumes model structure and latent dimension are fixed, and all latent variables are continuous.
Parameters: model (callable) – a Pyro model Reference:
 [1] Automatic Differentiation Variational Inference,
 Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, David M. Blei
Parameters:  model (callable) – A Pyro model.
 init_loc_fn (callable) – A persite initialization function. See Initialization section for available functions.

forward
(*args, **kwargs)[source]¶ An automatic guide with the same
*args, **kwargs
as the basemodel
.Returns: A dict mapping sample site name to sampled value. Return type: dict

median
(*args, **kwargs)[source]¶ Returns the posterior median value of each latent variable.
Returns: A dict mapping sample site name to median tensor. Return type: dict

quantiles
(quantiles, *args, **kwargs)[source]¶ Returns posterior quantiles each latent variable. Example:
print(guide.quantiles([0.05, 0.5, 0.95]))
Parameters: quantiles (torch.Tensor or 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
AutoMultivariateNormal¶

class
AutoMultivariateNormal
(model, init_loc_fn=<function init_to_median>, init_scale=0.1)[source]¶ Bases:
pyro.infer.autoguide.guides.AutoContinuous
This implementation of
AutoContinuous
uses a Cholesky factorization of a Multivariate Normal 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, ...)
By default the mean vector is initialized by
init_loc_fn()
and the Cholesky factor is initialized to the identity times a small factor.Parameters:  model (callable) – A generative model.
 init_loc_fn (callable) – A persite initialization function. See Initialization section for available functions.
 init_scale (float) – Initial scale for the standard deviation of each (unconstrained transformed) latent variable.
AutoDiagonalNormal¶

class
AutoDiagonalNormal
(model, init_loc_fn=<function init_to_median>, init_scale=0.1)[source]¶ Bases:
pyro.infer.autoguide.guides.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, ...)
By default the mean vector is initialized to zero and the scale is initialized to the identity times a small factor.
Parameters:  model (callable) – A generative model.
 init_loc_fn (callable) – A persite initialization function. See Initialization section for available functions.
 init_scale (float) – Initial scale for the standard deviation of each (unconstrained transformed) latent variable.
AutoLowRankMultivariateNormal¶

class
AutoLowRankMultivariateNormal
(model, init_loc_fn=<function init_to_median>, init_scale=0.1, rank=None)[source]¶ Bases:
pyro.infer.autoguide.guides.AutoContinuous
This implementation of
AutoContinuous
uses a low rank plus diagonal Multivariate Normal 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=10) svi = SVI(model, guide, ...)
By default the
cov_diag
is initialized to a small constant and thecov_factor
is initialized randomly such that on averagecov_factor.matmul(cov_factor.t())
has the same scale ascov_diag
.Parameters:  model (callable) – A generative model.
 rank (int or None) – The rank of the lowrank part of the covariance matrix.
Defaults to approximately
sqrt(latent dim)
.  init_loc_fn (callable) – A persite initialization function. See Initialization section for available functions.
 init_scale (float) – Approximate initial scale for the standard deviation of each (unconstrained transformed) latent variable.
AutoIAFNormal¶

class
AutoIAFNormal
(model, hidden_dim=None, init_loc_fn=<function init_to_median>)[source]¶ Bases:
pyro.infer.autoguide.guides.AutoContinuous
This implementation of
AutoContinuous
uses a Diagonal Normal distribution transformed via aAffineAutoregressive
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_dim=latent_dim) svi = SVI(model, guide, ...)
Parameters:  model (callable) – a generative model
 hidden_dim (int) – number of hidden dimensions in the IAF
 init_loc_fn (callable) – A persite initialization function. See Initialization section for available functions.
AutoLaplaceApproximation¶

class
AutoLaplaceApproximation
(model, init_loc_fn=<function init_to_median>)[source]¶ Bases:
pyro.infer.autoguide.guides.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:
delta_guide = AutoLaplaceApproximation(model) svi = SVI(model, delta_guide, ...) # ...then train the delta_guide... guide = delta_guide.laplace_approximation()
By default the mean vector is initialized to an empirical prior median.
Parameters:  model (callable) – a generative model
 init_loc_fn (callable) – A persite initialization function. See Initialization section for available functions.

laplace_approximation
(*args, **kwargs)[source]¶ Returns a
AutoMultivariateNormal
instance whose posterior’s loc and scale_tril are given by Laplace approximation.
AutoDiscreteParallel¶
Initialization¶
The pyro.infer.autoguide.initialization module contains initialization functions for automatic guides.
The standard interface for initialization is a function that inputs a Pyro
trace site
dict and returns an appropriately sized value
to serve
as an initial constrained value for a guide estimate.

init_to_feasible
(site)[source]¶ Initialize to an arbitrary feasible point, ignoring distribution parameters.

init_to_median
(site, num_samples=15)[source]¶ Initialize to the prior median; fallback to a feasible point if median is undefined.

class
InitMessenger
(init_fn)[source]¶ Bases:
pyro.poutine.messenger.Messenger
Initializes a site by replacing
.sample()
calls with values drawn from an initialization strategy. This is mainly for internal use by autoguide classes.Parameters: init_fn (callable) – An initialization function.
Reparameterizers¶
The pyro.infer.reparam
module contains reparameterization strategies for
the pyro.poutine.handlers.reparam()
effect. These are useful for altering
geometry of a poorlyconditioned parameter space to make the posterior better
shaped. These can be used with a variety of inference algorithms, e.g.
Auto*Normal
guides and MCMC.

class
Reparam
[source]¶ Base class for reparameterizers.

__call__
(name, fn, obs)[source]¶ Parameters:  name (str) – A sample site name.
 fn (TorchDistribution) – A distribution.
 obs (Tensor) – Observed value or None.
Returns: A pair (
new_fn
,value
).

LocScale Decentering¶

class
LocScaleReparam
(centered=None, shape_params=())[source]¶ Bases:
pyro.infer.reparam.reparam.Reparam
Generic decentering reparameterizer [1] for latent variables parameterized by
loc
andscale
(and possibly additionalshape_params
).This reparameterization works only for latent variables, not likelihoods.
 [1] Maria I. Gorinova, Dave Moore, Matthew D. Hoffman (2019)
 “Automatic Reparameterisation of Probabilistic Programs” https://arxiv.org/pdf/1906.03028.pdf
Parameters:  centered (float) – optional centered parameter. If None (default) learn
a persite perelement centering parameter in
[0,1]
. If 0, fully decenter the distribution; if 1, preserve the centered distribution unchanged.  shape_params (tuple or list) – list of additional parameter names to copy unchanged from the centered to decentered distribution.
Transformed Distributions¶

class
TransformReparam
[source]¶ Bases:
pyro.infer.reparam.reparam.Reparam
Reparameterizer for
pyro.distributions.torch.TransformedDistribution
latent variables.This is useful for transformed distributions with complex, geometrychanging transforms, where the posterior has simple shape in the space of
base_dist
.This reparameterization works only for latent variables, not likelihoods.
Discrete Cosine Transform¶

class
DiscreteCosineReparam
(dim=1)[source]¶ Bases:
pyro.infer.reparam.reparam.Reparam
Discrete Cosine reparamterizer, using a
DiscreteCosineTransform
.This is useful for sequential models where coupling along a timelike axis (e.g. a banded precision matrix) introduces longrange correlation. This reparameterizes to a frequencydomain represetation where posterior covariance should be closer to diagonal, thereby improving the accuracy of diagonal guides in SVI and improving the effectiveness of a diagonal mass matrix in HMC.
This reparameterization works only for latent variables, not likelihoods.
Parameters: dim (int) – Dimension along which to transform. Must be negative. This is an absolute dim counting from the right.
StudentT Distributions¶

class
StudentTReparam
[source]¶ Bases:
pyro.infer.reparam.reparam.Reparam
Auxiliary variable reparameterizer for
StudentT
random variables.This is useful in combination with
LinearHMMReparam
because it allows StudentT processes to be treated as conditionally Gaussian processes, permitting cheap inference viaGaussianHMM
.This reparameterizes a
StudentT
by introducing an auxiliaryGamma
variable conditioned on which the result isNormal
.
Stable Distributions¶

class
LatentStableReparam
[source]¶ Bases:
pyro.infer.reparam.reparam.Reparam
Auxiliary variable reparameterizer for
Stable
latent variables.This is useful in inference of latent
Stable
variables because thelog_prob()
is not implemented.This uses the ChambersMallowsStuck method [1], creating a pair of parameterfree auxiliary distributions (
Uniform(pi/2,pi/2)
andExponential(1)
) with welldefined.log_prob()
methods, thereby permitting use of reparameterized stable distributions in likelihoodbased inference algorithms like SVI and MCMC.This reparameterization works only for latent variables, not likelihoods. For likelihoodcompatible reparameterization see
SymmetricStableReparam
orStableReparam
. [1] J.P. Nolan (2017).
 Stable Distributions: Models for Heavy Tailed Data. http://fs2.american.edu/jpnolan/www/stable/chap1.pdf

class
SymmetricStableReparam
[source]¶ Bases:
pyro.infer.reparam.reparam.Reparam
Auxiliary variable reparameterizer for symmetric
Stable
random variables (i.e. those for whichskew=0
).This is useful in inference of symmetric
Stable
variables because thelog_prob()
is not implemented.This reparameterizes a symmetric
Stable
random variable as a totallyskewed (skew=1
)Stable
scale mixture ofNormal
random variables. See Proposition 3. of [1] (but note we differ sinceStable
uses Nolan’s continuous S0 parameterization). [1] Alvaro Cartea and Sam Howison (2009)
 “Option Pricing with LevyStable Processes” https://pdfs.semanticscholar.org/4d66/c91b136b2a38117dd16c2693679f5341c616.pdf

class
StableReparam
[source]¶ Bases:
pyro.infer.reparam.reparam.Reparam
Auxiliary variable reparameterizer for arbitrary
Stable
random variables.This is useful in inference of nonsymmetric
Stable
variables because thelog_prob()
is not implemented.This reparameterizes a
Stable
random variable as sum of two other stable random variables, one symmetric and the other totally skewed (applying Property 2.3.a of [1]). The totally skewed variable is sampled as inLatentStableReparam
, and the symmetric variable is decomposed as inSymmetricStableReparam
. [1] V. M. Zolotarev (1986)
 “Onedimensional stable distributions”
Hidden Markov Models¶

class
LinearHMMReparam
(init=None, trans=None, obs=None)[source]¶ Bases:
pyro.infer.reparam.reparam.Reparam
Auxiliary variable reparameterizer for
LinearHMM
random variables.This defers to component reparameterizers to create auxiliary random variables conditioned on which the process becomes a
GaussianHMM
. If theobservation_dist
is aTransformedDistribution
this reorders those transforms so that the result is aTransformedDistribution
ofGaussianHMM
.This is useful for training the parameters of a
LinearHMM
distribution, whoselog_prob()
method is undefined. To perform inference in the presence of nonGaussian factors such asStable()
,StudentT()
orLogNormal()
, configure withStudentTReparam
,StableReparam
,SymmetricStableReparam
, etc. component reparameterizers forinit
,trans
, andscale
. For example:hmm = LinearHMM( init_dist=Stable(1,0,1,0).expand([2]).to_event(1), trans_matrix=torch.eye(2), trans_dist=MultivariateNormal(torch.zeros(2), torch.eye(2)), obs_matrix=torch.eye(2), obs_dist=TransformedDistribution( Stable(1.5,0.5,1.0).expand([2]).to_event(1), ExpTransform())) rep = LinearHMMReparam(init=SymmetricStableReparam(), obs=StableReparam()) with poutine.reparam(config={"hmm": rep}): pyro.sample("hmm", hmm, obs=data)
Parameters:
Neural Transport¶

class
NeuTraReparam
(guide)[source]¶ Bases:
pyro.infer.reparam.reparam.Reparam
Neural Transport reparameterizer [1] of multiple latent variables.
This uses a trained
AutoContinuous
guide to alter the geometry of a model, typically for use e.g. in MCMC. Example usage:# Step 1. Train a guide guide = AutoIAFNormal(model) svi = SVI(model, guide, ...) # ...train the guide... # Step 2. Use trained guide in NeuTra MCMC neutra = NeuTraReparam(guide) model = poutine.reparam(model, config=lambda _: neutra) nuts = NUTS(model) # ...now use the model in HMC or NUTS...
This reparameterization works only for latent variables, not likelihoods. Note that all sites must share a single common
NeuTraReparam
instance, and that the model must have static structure. [1] Hoffman, M. et al. (2019)
 “NeuTralizing Bad Geometry in Hamiltonian Monte Carlo Using Neural Transport” https://arxiv.org/abs/1903.03704
Parameters: guide (AutoContinuous) – A trained guide. 
transform_sample
(latent)[source]¶ Given latent samples from the warped posterior (with a possible batch dimension), return a dict of samples from the latent sites in the model.
Parameters: latent – sample from the warped posterior (possibly batched). Note that the batch dimension must not collide with plate dimensions in the model, i.e. any batch dims d <  max_plate_nesting. Returns: a dict of samples keyed by latent sites in the model. Return type: dict
Distributions¶
PyTorch Distributions¶
Most distributions in Pyro are thin wrappers around PyTorch distributions.
For details on the PyTorch distribution interface, see
torch.distributions.distribution.Distribution
.
For differences between the Pyro and PyTorch interfaces, see
TorchDistributionMixin
.
Bernoulli¶

class
Bernoulli
(probs=None, logits=None, validate_args=None)¶ Wraps
torch.distributions.bernoulli.Bernoulli
withTorchDistributionMixin
.
Beta¶

class
Beta
(concentration1, concentration0, validate_args=None)¶ Wraps
torch.distributions.beta.Beta
withTorchDistributionMixin
.
Binomial¶

class
Binomial
(total_count=1, probs=None, logits=None, validate_args=None)¶ Wraps
torch.distributions.binomial.Binomial
withTorchDistributionMixin
.
Categorical¶

class
Categorical
(probs=None, logits=None, validate_args=None)[source]¶ Wraps
torch.distributions.categorical.Categorical
withTorchDistributionMixin
.
Cauchy¶

class
Cauchy
(loc, scale, validate_args=None)¶ Wraps
torch.distributions.cauchy.Cauchy
withTorchDistributionMixin
.
Chi2¶

class
Chi2
(df, validate_args=None)¶ Wraps
torch.distributions.chi2.Chi2
withTorchDistributionMixin
.
Dirichlet¶

class
Dirichlet
(concentration, validate_args=None)¶ Wraps
torch.distributions.dirichlet.Dirichlet
withTorchDistributionMixin
.
Exponential¶

class
Exponential
(rate, validate_args=None)¶ Wraps
torch.distributions.exponential.Exponential
withTorchDistributionMixin
.
ExponentialFamily¶

class
ExponentialFamily
(batch_shape=torch.Size([]), event_shape=torch.Size([]), validate_args=None)¶ Wraps
torch.distributions.exp_family.ExponentialFamily
withTorchDistributionMixin
.
FisherSnedecor¶

class
FisherSnedecor
(df1, df2, validate_args=None)¶ Wraps
torch.distributions.fishersnedecor.FisherSnedecor
withTorchDistributionMixin
.
Gamma¶

class
Gamma
(concentration, rate, validate_args=None)¶ Wraps
torch.distributions.gamma.Gamma
withTorchDistributionMixin
.
Geometric¶

class
Geometric
(probs=None, logits=None, validate_args=None)¶ Wraps
torch.distributions.geometric.Geometric
withTorchDistributionMixin
.
Gumbel¶

class
Gumbel
(loc, scale, validate_args=None)¶ Wraps
torch.distributions.gumbel.Gumbel
withTorchDistributionMixin
.
HalfCauchy¶

class
HalfCauchy
(scale, validate_args=None)¶ Wraps
torch.distributions.half_cauchy.HalfCauchy
withTorchDistributionMixin
.
HalfNormal¶

class
HalfNormal
(scale, validate_args=None)¶ Wraps
torch.distributions.half_normal.HalfNormal
withTorchDistributionMixin
.
Independent¶

class
Independent
(base_distribution, reinterpreted_batch_ndims, validate_args=None)[source]¶ Wraps
torch.distributions.independent.Independent
withTorchDistributionMixin
.
Laplace¶

class
Laplace
(loc, scale, validate_args=None)¶ Wraps
torch.distributions.laplace.Laplace
withTorchDistributionMixin
.
LogNormal¶

class
LogNormal
(loc, scale, validate_args=None)¶ Wraps
torch.distributions.log_normal.LogNormal
withTorchDistributionMixin
.
LogisticNormal¶

class
LogisticNormal
(loc, scale, validate_args=None)¶ Wraps
torch.distributions.logistic_normal.LogisticNormal
withTorchDistributionMixin
.
LowRankMultivariateNormal¶

class
LowRankMultivariateNormal
(loc, cov_factor, cov_diag, validate_args=None)¶ Wraps
torch.distributions.lowrank_multivariate_normal.LowRankMultivariateNormal
withTorchDistributionMixin
.
Multinomial¶

class
Multinomial
(total_count=1, probs=None, logits=None, validate_args=None)¶ Wraps
torch.distributions.multinomial.Multinomial
withTorchDistributionMixin
.
MultivariateNormal¶

class
MultivariateNormal
(loc, covariance_matrix=None, precision_matrix=None, scale_tril=None, validate_args=None)[source]¶ Wraps
torch.distributions.multivariate_normal.MultivariateNormal
withTorchDistributionMixin
.
NegativeBinomial¶

class
NegativeBinomial
(total_count, probs=None, logits=None, validate_args=None)¶ Wraps
torch.distributions.negative_binomial.NegativeBinomial
withTorchDistributionMixin
.
Normal¶

class
Normal
(loc, scale, validate_args=None)¶ Wraps
torch.distributions.normal.Normal
withTorchDistributionMixin
.
OneHotCategorical¶

class
OneHotCategorical
(probs=None, logits=None, validate_args=None)¶ Wraps
torch.distributions.one_hot_categorical.OneHotCategorical
withTorchDistributionMixin
.
Pareto¶

class
Pareto
(scale, alpha, validate_args=None)¶ Wraps
torch.distributions.pareto.Pareto
withTorchDistributionMixin
.
Poisson¶

class
Poisson
(rate, validate_args=None)¶ Wraps
torch.distributions.poisson.Poisson
withTorchDistributionMixin
.
RelaxedBernoulli¶

class
RelaxedBernoulli
(temperature, probs=None, logits=None, validate_args=None)¶ Wraps
torch.distributions.relaxed_bernoulli.RelaxedBernoulli
withTorchDistributionMixin
.
RelaxedOneHotCategorical¶

class
RelaxedOneHotCategorical
(temperature, probs=None, logits=None, validate_args=None)¶ Wraps
torch.distributions.relaxed_categorical.RelaxedOneHotCategorical
withTorchDistributionMixin
.
StudentT¶

class
StudentT
(df, loc=0.0, scale=1.0, validate_args=None)¶ Wraps
torch.distributions.studentT.StudentT
withTorchDistributionMixin
.
TransformedDistribution¶

class
TransformedDistribution
(base_distribution, transforms, validate_args=None)¶ Wraps
torch.distributions.transformed_distribution.TransformedDistribution
withTorchDistributionMixin
.
Uniform¶

class
Uniform
(low, high, validate_args=None)¶ Wraps
torch.distributions.uniform.Uniform
withTorchDistributionMixin
.
Weibull¶

class
Weibull
(scale, concentration, validate_args=None)¶ Wraps
torch.distributions.weibull.Weibull
withTorchDistributionMixin
.
Pyro Distributions¶
Abstract Distribution¶

class
Distribution
[source]¶ Bases:
object
Base class for parameterized probability distributions.
Distributions in Pyro are stochastic function objects with
sample()
andlog_prob()
methods. Distribution are stochastic functions with fixed parameters:d = dist.Bernoulli(param) x = d() # Draws a random sample. p = d.log_prob(x) # Evaluates log probability of x.
Implementing New Distributions:
Derived classes must implement the methods:
sample()
,log_prob()
.Examples:
Take a look at the examples to see how they interact with inference algorithms.

__call__
(*args, **kwargs)[source]¶ Samples a random value (just an alias for
.sample(*args, **kwargs)
).For tensor distributions, the returned tensor should have the same
.shape
as the parameters.Returns: A random value. Return type: torch.Tensor

enumerate_support
(expand=True)[source]¶ Returns a representation of the parametrized distribution’s support, along the first dimension. This is implemented only by discrete distributions.
Note that this returns support values of all the batched RVs in lockstep, rather than the full cartesian product.
Parameters: expand (bool) – whether to expand the result to a tensor of shape (n,) + batch_shape + event_shape
. If false, the return value has unexpanded shape(n,) + (1,)*len(batch_shape) + event_shape
which can be broadcasted to the full shape.Returns: An iterator over the distribution’s discrete support. Return type: iterator

has_enumerate_support
= False¶

has_rsample
= False¶

log_prob
(x, *args, **kwargs)[source]¶ Evaluates log probability densities for each of a batch of samples.
Parameters: x (torch.Tensor) – A single value or a batch of values batched along axis 0. Returns: log probability densities as a onedimensional Tensor
with same batch size as value and params. The shape of the result should beself.batch_size
.Return type: torch.Tensor

sample
(*args, **kwargs)[source]¶ Samples a random value.
For tensor distributions, the returned tensor should have the same
.shape
as the parameters, unless otherwise noted.Parameters: sample_shape (torch.Size) – the size of the iid batch to be drawn from the distribution. Returns: A random value or batch of random values (if parameters are batched). The shape of the result should be self.shape()
.Return type: torch.Tensor

score_parts
(x, *args, **kwargs)[source]¶ Computes ingredients for stochastic gradient estimators of ELBO.
The default implementation is correct both for nonreparameterized and for fully reparameterized distributions. Partially reparameterized distributions should override this method to compute correct .score_function and .entropy_term parts.
Parameters: x (torch.Tensor) – A single value or batch of values. Returns: A ScoreParts object containing parts of the ELBO estimator. Return type: ScoreParts

TorchDistributionMixin¶

class
TorchDistributionMixin
[source]¶ Bases:
pyro.distributions.distribution.Distribution
Mixin to provide Pyro compatibility for PyTorch distributions.
You should instead use TorchDistribution for new distribution classes.
This is mainly useful for wrapping existing PyTorch distributions for use in Pyro. Derived classes must first inherit from
torch.distributions.distribution.Distribution
and then inherit fromTorchDistributionMixin
.
__call__
(sample_shape=torch.Size([]))[source]¶ Samples a random value.
This is reparameterized whenever possible, calling
rsample()
for reparameterized distributions andsample()
for nonreparameterized distributions.Parameters: sample_shape (torch.Size) – the size of the iid batch to be drawn from the distribution. Returns: A random value or batch of random values (if parameters are batched). The shape of the result should be self.shape(). Return type: torch.Tensor

shape
(sample_shape=torch.Size([]))[source]¶ The tensor shape of samples from this distribution.
Samples are of shape:
d.shape(sample_shape) == sample_shape + d.batch_shape + d.event_shape
Parameters: sample_shape (torch.Size) – the size of the iid batch to be drawn from the distribution. Returns: Tensor shape of samples. Return type: torch.Size

expand
(batch_shape, _instance=None)[source]¶ Returns a new
ExpandedDistribution
instance with batch dimensions expanded to batch_shape.Parameters:  batch_shape (tuple) – batch shape to expand to.
 _instance – unused argument for compatibility with
torch.distributions.Distribution.expand()
Returns: an instance of ExpandedDistribution.
Return type: ExpandedDistribution

expand_by
(sample_shape)[source]¶ Expands a distribution by adding
sample_shape
to the left side of itsbatch_shape
.To expand internal dims of
self.batch_shape
from 1 to something larger, useexpand()
instead.Parameters: sample_shape (torch.Size) – The size of the iid batch to be drawn from the distribution. Returns: An expanded version of this distribution. Return type: ExpandedDistribution

to_event
(reinterpreted_batch_ndims=None)[source]¶ Reinterprets the
n
rightmost dimensions of this distributionsbatch_shape
as event dims, adding them to the left side ofevent_shape
.Example:
>>> [d1.batch_shape, d1.event_shape] [torch.Size([2, 3]), torch.Size([4, 5])] >>> d2 = d1.to_event(1) >>> [d2.batch_shape, d2.event_shape] [torch.Size([2]), torch.Size([3, 4, 5])] >>> d3 = d1.to_event(2) >>> [d3.batch_shape, d3.event_shape] [torch.Size([]), torch.Size([2, 3, 4, 5])]
Parameters: reinterpreted_batch_ndims (int) – The number of batch dimensions to reinterpret as event dimensions. May be negative to remove dimensions from an pyro.distributions.torch.Independent
. If None, convert all dimensions to event dimensions.Returns: A reshaped version of this distribution. Return type: pyro.distributions.torch.Independent

mask
(mask)[source]¶ Masks a distribution by a boolean or booleanvalued tensor that is broadcastable to the distributions
batch_shape
.Parameters: mask (bool or torch.Tensor) – A boolean or boolean valued tensor. Returns: A masked copy of this distribution. Return type: MaskedDistribution

TorchDistribution¶

class
TorchDistribution
(batch_shape=torch.Size([]), event_shape=torch.Size([]), validate_args=None)[source]¶ Bases:
torch.distributions.distribution.Distribution
,pyro.distributions.torch_distribution.TorchDistributionMixin
Base class for PyTorchcompatible distributions with Pyro support.
This should be the base class for almost all new Pyro distributions.
Note
Parameters and data should be of type
Tensor
and all methods return typeTensor
unless otherwise noted.Tensor Shapes:
TorchDistributions provide a method
.shape()
for the tensor shape of samples:x = d.sample(sample_shape) assert x.shape == d.shape(sample_shape)
Pyro follows the same distribution shape semantics as PyTorch. It distinguishes between three different roles for tensor shapes of samples:
 sample shape corresponds to the shape of the iid samples drawn from the distribution. This is taken as an argument by the distribution’s sample method.
 batch shape corresponds to nonidentical (independent) parameterizations of the distribution, inferred from the distribution’s parameter shapes. This is fixed for a distribution instance.
 event shape corresponds to the event dimensions of the distribution, which is fixed for a distribution class. These are collapsed when we try to score a sample from the distribution via d.log_prob(x).
These shapes are related by the equation:
assert d.shape(sample_shape) == sample_shape + d.batch_shape + d.event_shape
Distributions provide a vectorized
log_prob()
method that evaluates the log probability density of each event in a batch independently, returning a tensor of shapesample_shape + d.batch_shape
:x = d.sample(sample_shape) assert x.shape == d.shape(sample_shape) log_p = d.log_prob(x) assert log_p.shape == sample_shape + d.batch_shape
Implementing New Distributions:
Derived classes must implement the methods
sample()
(orrsample()
if.has_rsample == True
) andlog_prob()
, and must implement the propertiesbatch_shape
, andevent_shape
. Discrete classes may also implement theenumerate_support()
method to improve gradient estimates and set.has_enumerate_support = True
.
expand
(batch_shape, _instance=None)¶ Returns a new
ExpandedDistribution
instance with batch dimensions expanded to batch_shape.Parameters:  batch_shape (tuple) – batch shape to expand to.
 _instance – unused argument for compatibility with
torch.distributions.Distribution.expand()
Returns: an instance of ExpandedDistribution.
Return type: ExpandedDistribution
AVFMultivariateNormal¶

class
AVFMultivariateNormal
(loc, scale_tril, control_var)[source]¶ Bases:
pyro.distributions.torch.MultivariateNormal
Multivariate normal (Gaussian) distribution with transport equation inspired control variates (adaptive velocity fields).
A distribution over vectors in which all the elements have a joint Gaussian density.
Parameters:  loc (torch.Tensor) – Ddimensional mean vector.
 scale_tril (torch.Tensor) – Cholesky of Covariance matrix; D x D matrix.
 control_var (torch.Tensor) – 2 x L x D tensor that parameterizes the control variate; L is an arbitrary positive integer. This parameter needs to be learned (i.e. adapted) to achieve lower variance gradients. In a typical use case this parameter will be adapted concurrently with the loc and scale_tril that define the distribution.
Example usage:
control_var = torch.tensor(0.1 * torch.ones(2, 1, D), requires_grad=True) opt_cv = torch.optim.Adam([control_var], lr=0.1, betas=(0.5, 0.999)) for _ in range(1000): d = AVFMultivariateNormal(loc, scale_tril, control_var) z = d.rsample() cost = torch.pow(z, 2.0).sum() cost.backward() opt_cv.step() opt_cv.zero_grad()

arg_constraints
= {'control_var': Real(), 'loc': Real(), 'scale_tril': LowerTriangular()}¶
BetaBinomial¶

class
BetaBinomial
(concentration1, concentration0, total_count=1, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Compound distribution comprising of a betabinomial pair. The probability of success (
probs
for theBinomial
distribution) is unknown and randomly drawn from aBeta
distribution prior to a certain number of Bernoulli trials given bytotal_count
.Parameters: 
arg_constraints
= {'concentration0': GreaterThan(lower_bound=0.0), 'concentration1': GreaterThan(lower_bound=0.0), 'total_count': IntegerGreaterThan(lower_bound=0)}¶

concentration0
¶

concentration1
¶

has_enumerate_support
= True¶

mean
¶

support
¶

variance
¶

ConditionalDistribution¶
ConditionalTransformedDistribution¶
Delta¶

class
Delta
(v, log_density=0.0, event_dim=0, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Degenerate discrete distribution (a single point).
Discrete distribution that assigns probability one to the single element in its support. Delta distribution parameterized by a random choice should not be used with MCMC based inference, as doing so produces incorrect results.
Parameters:  v (torch.Tensor) – The single support element.
 log_density (torch.Tensor) – An optional density for this Delta. This
is useful to keep the class of
Delta
distributions closed under differentiable transformation.  event_dim (int) – Optional event dimension, defaults to zero.

arg_constraints
= {'log_density': Real(), 'v': Real()}¶

has_rsample
= True¶

mean
¶

support
= Real()¶

variance
¶
DirichletMultinomial¶

class
DirichletMultinomial
(concentration, total_count=1, is_sparse=False, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Compound distribution comprising of a dirichletmultinomial pair. The probability of classes (
probs
for theMultinomial
distribution) is unknown and randomly drawn from aDirichlet
distribution prior to a certain number of Categorical trials given bytotal_count
.Parameters:  or torch.Tensor concentration (float) – concentration parameter (alpha) for the Dirichlet distribution.
 or torch.Tensor total_count (int) – number of Categorical trials.
 is_sparse (bool) – Whether to assume value is mostly zero when computing
log_prob()
, which can speed up computation when data is sparse.

arg_constraints
= {'concentration': GreaterThan(lower_bound=0.0), 'total_count': IntegerGreaterThan(lower_bound=0)}¶

concentration
¶

mean
¶

support
¶

variance
¶
DiscreteHMM¶

class
DiscreteHMM
(initial_logits, transition_logits, observation_dist, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Hidden Markov Model with discrete latent state and arbitrary observation distribution. This uses [1] to parallelize over time, achieving O(log(time)) parallel complexity.
The event_shape of this distribution includes time on the left:
event_shape = (num_steps,) + observation_dist.event_shape
This distribution supports any combination of homogeneous/heterogeneous time dependency of
transition_logits
andobservation_dist
. However, because time is included in this distribution’s event_shape, the homogeneous+homogeneous case will have a broadcastable event_shape withnum_steps = 1
, allowinglog_prob()
to work with arbitrary length data:# homogeneous + homogeneous case: event_shape = (1,) + observation_dist.event_shape
References:
 [1] Simo Sarkka, Angel F. GarciaFernandez (2019)
 “Temporal Parallelization of Bayesian Filters and Smoothers” https://arxiv.org/pdf/1905.13002.pdf
Parameters:  initial_logits (Tensor) – A logits tensor for an initial
categorical distribution over latent states. Should have rightmost size
state_dim
and be broadcastable tobatch_shape + (state_dim,)
.  transition_logits (Tensor) – A logits tensor for transition
conditional distributions between latent states. Should have rightmost
shape
(state_dim, state_dim)
(old, new), and be broadcastable tobatch_shape + (num_steps, state_dim, state_dim)
.  observation_dist (Distribution) – A conditional
distribution of observed data conditioned on latent state. The
.batch_shape
should have rightmost sizestate_dim
and be broadcastable tobatch_shape + (num_steps, state_dim)
. The.event_shape
may be arbitrary.

arg_constraints
= {'initial_logits': Real(), 'transition_logits': Real()}¶

filter
(value)[source]¶ Compute posterior over final state given a sequence of observations.
Parameters: value (Tensor) – A sequence of observations. Returns: A posterior distribution over latent states at the final time step. result.logits
can then be used asinitial_logits
in a sequential Pyro model for prediction.Return type: Categorical
EmpiricalDistribution¶

class
Empirical
(samples, log_weights, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Empirical distribution associated with the sampled data. Note that the shape requirement for log_weights is that its shape must match the leftmost shape of samples. Samples are aggregated along the
aggregation_dim
, which is the rightmost dim of log_weights.Example:
>>> emp_dist = Empirical(torch.randn(2, 3, 10), torch.ones(2, 3)) >>> emp_dist.batch_shape torch.Size([2]) >>> emp_dist.event_shape torch.Size([10])
>>> single_sample = emp_dist.sample() >>> single_sample.shape torch.Size([2, 10]) >>> batch_sample = emp_dist.sample((100,)) >>> batch_sample.shape torch.Size([100, 2, 10])
>>> emp_dist.log_prob(single_sample).shape torch.Size([2]) >>> # Vectorized samples cannot be scored by log_prob. >>> with pyro.validation_enabled(): ... emp_dist.log_prob(batch_sample).shape Traceback (most recent call last): ... ValueError: ``value.shape`` must be torch.Size([2, 10])
Parameters:  samples (torch.Tensor) – samples from the empirical distribution.
 log_weights (torch.Tensor) – log weights (optional) corresponding to the samples.

arg_constraints
= {}¶

enumerate_support
(expand=True)[source]¶ See
pyro.distributions.torch_distribution.TorchDistribution.enumerate_support()

event_shape
¶ See
pyro.distributions.torch_distribution.TorchDistribution.event_shape()

has_enumerate_support
= True¶

log_prob
(value)[source]¶ Returns the log of the probability mass function evaluated at
value
. Note that this currently only supports scoring values with emptysample_shape
.Parameters: value (torch.Tensor) – scalar or tensor value to be scored.

log_weights
¶

mean
¶ See
pyro.distributions.torch_distribution.TorchDistribution.mean()

sample
(sample_shape=torch.Size([]))[source]¶ See
pyro.distributions.torch_distribution.TorchDistribution.sample()

sample_size
¶ Number of samples that constitute the empirical distribution.
Return int: number of samples collected.

support
= Real()¶

variance
¶ See
pyro.distributions.torch_distribution.TorchDistribution.variance()
FoldedDistribution¶

class
FoldedDistribution
(base_dist, validate_args=None)[source]¶ Bases:
pyro.distributions.torch.TransformedDistribution
Equivalent to
TransformedDistribution(base_dist, AbsTransform())
, but additionally supportslog_prob()
.Parameters: base_dist (Distribution) – The distribution to reflect. 
support
= GreaterThan(lower_bound=0.0)¶

GammaGaussianHMM¶

class
GammaGaussianHMM
(scale_dist, initial_dist, transition_matrix, transition_dist, observation_matrix, observation_dist, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Hidden Markov Model with the joint distribution of initial state, hidden state, and observed state is a
MultivariateStudentT
distribution along the line of references [2] and [3]. This adapts [1] to parallelize over time to achieve O(log(time)) parallel complexity.This GammaGaussianHMM class corresponds to the generative model:
s = Gamma(df/2, df/2).sample() z = scale(initial_dist, s).sample() x = [] for t in range(num_events): z = z @ transition_matrix + scale(transition_dist, s).sample() x.append(z @ observation_matrix + scale(observation_dist, s).sample())
where scale(mvn(loc, precision), s) := mvn(loc, s * precision).
The event_shape of this distribution includes time on the left:
event_shape = (num_steps,) + observation_dist.event_shape
This distribution supports any combination of homogeneous/heterogeneous time dependency of
transition_dist
andobservation_dist
. However, because time is included in this distribution’s event_shape, the homogeneous+homogeneous case will have a broadcastable event_shape withnum_steps = 1
, allowinglog_prob()
to work with arbitrary length data:event_shape = (1, obs_dim) # homogeneous + homogeneous case
References:
 [1] Simo Sarkka, Angel F. GarciaFernandez (2019)
 “Temporal Parallelization of Bayesian Filters and Smoothers” https://arxiv.org/pdf/1905.13002.pdf
 [2] F. J. Giron and J. C. Rojano (1994)
 “Bayesian Kalman filtering with elliptically contoured errors”
 [3] Filip Tronarp, Toni Karvonen, and Simo Sarkka (2019)
 “Student’s tfilters for noise scale estimation” https://users.aalto.fi/~ssarkka/pub/SPL2019.pdf
Variables: Parameters:  scale_dist (Gamma) – Prior of the mixing distribution.
 initial_dist (MultivariateNormal) – A distribution with unit scale mixing
over initial states. This should have batch_shape broadcastable to
self.batch_shape
. This should have event_shape(hidden_dim,)
.  transition_matrix (Tensor) – A linear transformation of hidden
state. This should have shape broadcastable to
self.batch_shape + (num_steps, hidden_dim, hidden_dim)
where the rightmost dims are ordered(old, new)
.  transition_dist (MultivariateNormal) – A process noise distribution
with unit scale mixing. This should have batch_shape broadcastable to
self.batch_shape + (num_steps,)
. This should have event_shape(hidden_dim,)
.  observation_matrix (Tensor) – A linear transformation from hidden
to observed state. This should have shape broadcastable to
self.batch_shape + (num_steps, hidden_dim, obs_dim)
.  observation_dist (MultivariateNormal) – An observation noise distribution
with unit scale mixing. This should have batch_shape broadcastable to
self.batch_shape + (num_steps,)
. This should have event_shape(obs_dim,)
.

arg_constraints
= {}¶

filter
(value)[source]¶ Compute posteriors over the multiplier and the final state given a sequence of observations. The posterior is a pair of Gamma and MultivariateNormal distributions (i.e. a GammaGaussian instance).
Parameters: value (Tensor) – A sequence of observations. Returns: A pair of posterior distributions over the mixing and the latent state at the final time step. Return type: a tuple of ~pyro.distributions.Gamma and ~pyro.distributions.MultivariateNormal
GammaPoisson¶

class
GammaPoisson
(concentration, rate, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Compound distribution comprising of a gammapoisson pair, also referred to as a gammapoisson mixture. The
rate
parameter for thePoisson
distribution is unknown and randomly drawn from aGamma
distribution.Note
This can be treated as an alternate parametrization of the
NegativeBinomial
(total_count
,probs
) distribution, with concentration = total_count and rate = (1  probs) / probs.Parameters: 
arg_constraints
= {'concentration': GreaterThan(lower_bound=0.0), 'rate': GreaterThan(lower_bound=0.0)}¶

concentration
¶

mean
¶

rate
¶

support
= IntegerGreaterThan(lower_bound=0)¶

variance
¶

GaussianHMM¶

class
GaussianHMM
(initial_dist, transition_matrix, transition_dist, observation_matrix, observation_dist, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Hidden Markov Model with Gaussians for initial, transition, and observation distributions. This adapts [1] to parallelize over time to achieve O(log(time)) parallel complexity, however it differs in that it tracks the log normalizer to ensure
log_prob()
is differentiable.This corresponds to the generative model:
z = initial_distribution.sample() x = [] for t in range(num_events): z = z @ transition_matrix + transition_dist.sample() x.append(z @ observation_matrix + observation_dist.sample())
The event_shape of this distribution includes time on the left:
event_shape = (num_steps,) + observation_dist.event_shape
This distribution supports any combination of homogeneous/heterogeneous time dependency of
transition_dist
andobservation_dist
. However, because time is included in this distribution’s event_shape, the homogeneous+homogeneous case will have a broadcastable event_shape withnum_steps = 1
, allowinglog_prob()
to work with arbitrary length data:event_shape = (1, obs_dim) # homogeneous + homogeneous case
References:
 [1] Simo Sarkka, Angel F. GarciaFernandez (2019)
 “Temporal Parallelization of Bayesian Filters and Smoothers” https://arxiv.org/pdf/1905.13002.pdf
Variables: Parameters:  initial_dist (MultivariateNormal) – A distribution
over initial states. This should have batch_shape broadcastable to
self.batch_shape
. This should have event_shape(hidden_dim,)
.  transition_matrix (Tensor) – A linear transformation of hidden
state. This should have shape broadcastable to
self.batch_shape + (num_steps, hidden_dim, hidden_dim)
where the rightmost dims are ordered(old, new)
.  transition_dist (MultivariateNormal) – A process
noise distribution. This should have batch_shape broadcastable to
self.batch_shape + (num_steps,)
. This should have event_shape(hidden_dim,)
.  observation_matrix (Tensor) – A linear transformation from hidden
to observed state. This should have shape broadcastable to
self.batch_shape + (num_steps, hidden_dim, obs_dim)
.  observation_dist (MultivariateNormal or
Normal) – An observation noise distribution. This should
have batch_shape broadcastable to
self.batch_shape + (num_steps,)
. This should have event_shape(obs_dim,)
.

arg_constraints
= {}¶

filter
(value)[source]¶ Compute posterior over final state given a sequence of observations.
Parameters: value (Tensor) – A sequence of observations. Returns: A posterior distribution over latent states at the final time step. result
can then be used asinitial_dist
in a sequential Pyro model for prediction.Return type: MultivariateNormal
GaussianMRF¶

class
GaussianMRF
(initial_dist, transition_dist, observation_dist, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Temporal Markov Random Field with Gaussian factors for initial, transition, and observation distributions. This adapts [1] to parallelize over time to achieve O(log(time)) parallel complexity, however it differs in that it tracks the log normalizer to ensure
log_prob()
is differentiable.The event_shape of this distribution includes time on the left:
event_shape = (num_steps,) + observation_dist.event_shape
This distribution supports any combination of homogeneous/heterogeneous time dependency of
transition_dist
andobservation_dist
. However, because time is included in this distribution’s event_shape, the homogeneous+homogeneous case will have a broadcastable event_shape withnum_steps = 1
, allowinglog_prob()
to work with arbitrary length data:event_shape = (1, obs_dim) # homogeneous + homogeneous case
References:
 [1] Simo Sarkka, Angel F. GarciaFernandez (2019)
 “Temporal Parallelization of Bayesian Filters and Smoothers” https://arxiv.org/pdf/1905.13002.pdf
Variables: Parameters:  initial_dist (MultivariateNormal) – A distribution
over initial states. This should have batch_shape broadcastable to
self.batch_shape
. This should have event_shape(hidden_dim,)
.  transition_dist (MultivariateNormal) – A joint
distribution factor over a pair of successive time steps. This should
have batch_shape broadcastable to
self.batch_shape + (num_steps,)
. This should have event_shape(hidden_dim + hidden_dim,)
(old+new).  observation_dist (MultivariateNormal) – A joint
distribution factor over a hidden and an observed state. This should
have batch_shape broadcastable to
self.batch_shape + (num_steps,)
. This should have event_shape(hidden_dim + obs_dim,)
.

arg_constraints
= {}¶
GaussianScaleMixture¶

class
GaussianScaleMixture
(coord_scale, component_logits, component_scale)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Mixture of Normal distributions with zero mean and diagonal covariance matrices.
That is, this distribution is a mixture with K components, where each component distribution is a Ddimensional Normal distribution with zero mean and a Ddimensional diagonal covariance matrix. The K different covariance matrices are controlled by the parameters coord_scale and component_scale. That is, the covariance matrix of the k’th component is given by
Sigma_ii = (component_scale_k * coord_scale_i) ** 2 (i = 1, …, D)
where component_scale_k is a positive scale factor and coord_scale_i are positive scale parameters shared between all K components. The mixture weights are controlled by a Kdimensional vector of softmax logits, component_logits. This distribution implements pathwise derivatives for samples from the distribution. This distribution does not currently support batched parameters.
See reference [1] for details on the implementations of the pathwise derivative. Please consider citing this reference if you use the pathwise derivative in your research.
[1] Pathwise Derivatives for Multivariate Distributions, Martin Jankowiak & Theofanis Karaletsos. arXiv:1806.01856
Note that this distribution supports both even and odd dimensions, but the former should be more a bit higher precision, since it doesn’t use any erfs in the backward call. Also note that this distribution does not support D = 1.
Parameters:  coord_scale (torch.tensor) – Ddimensional vector of scales
 component_logits (torch.tensor) – Kdimensional vector of logits
 component_scale (torch.tensor) – Kdimensional vector of scale multipliers

arg_constraints
= {'component_logits': Real(), 'component_scale': GreaterThan(lower_bound=0.0), 'coord_scale': GreaterThan(lower_bound=0.0)}¶

has_rsample
= True¶
InverseGamma¶

class
InverseGamma
(concentration, rate, validate_args=None)[source]¶ Bases:
pyro.distributions.torch.TransformedDistribution
Creates an inversegamma distribution parameterized by concentration and rate.
X ~ Gamma(concentration, rate) Y = 1/X ~ InverseGamma(concentration, rate)Parameters:  concentration (torch.Tensor) – the concentration parameter (i.e. alpha).
 rate (torch.Tensor) – the rate parameter (i.e. beta).

arg_constraints
= {'concentration': GreaterThan(lower_bound=0.0), 'rate': GreaterThan(lower_bound=0.0)}¶

concentration
¶

has_rsample
= True¶

rate
¶

support
= GreaterThan(lower_bound=0.0)¶
LinearHMM¶

class
LinearHMM
(initial_dist, transition_matrix, transition_dist, observation_matrix, observation_dist, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Hidden Markov Model with linear dynamics and observations and arbitrary noise for initial, transition, and observation distributions. Each of those distributions can be e.g. :class:`~pyro.distributions.MultivariateNormal or :class:`~pyro.distributions.Independent of :class:`~pyro.distributions.Normal, :class:`~pyro.distributions.StudentT, or ~pyro.distributions.Stable . Additionally the observation distribution may be constrained, e.g. ~pyro.distributions.LogNormal
This corresponds to the generative model:
z = initial_distribution.sample() x = [] for t in range(num_events): z = z @ transition_matrix + transition_dist.sample() y = z @ observation_matrix + obs_base_dist.sample() x.append(obs_transform(y))
where
observation_dist
is split intoobs_base_dist
and an optionalobs_transform
(defaulting to the identity).This implements a reparameterized
rsample()
method but does not implement alog_prob()
method. Derived classes may implementlog_prob()
.Inference without
log_prob()
can be performed using either reparameterization withLinearHMMReparam
or likelihoodfree algorithms such asEnergyDistance
. Note that while stable processes generally require a common shared stability parameter \(\alpha\) , this distribution and the above inference algorithms allow heterogeneous stability parameters.The event_shape of this distribution includes time on the left:
event_shape = (num_steps,) + observation_dist.event_shape
This distribution supports any combination of homogeneous/heterogeneous time dependency of
transition_dist
andobservation_dist
. However at least one of the distributions or matrices must be expanded to contain the time dimension.Variables: Parameters:  initial_dist – A distribution over initial states. This should have
batch_shape broadcastable to
self.batch_shape
. This should have event_shape(hidden_dim,)
.  transition_matrix (Tensor) – A linear transformation of hidden
state. This should have shape broadcastable to
self.batch_shape + (num_steps, hidden_dim, hidden_dim)
where the rightmost dims are ordered(old, new)
.  transition_dist – A distribution over process noise. This should have
batch_shape broadcastable to
self.batch_shape + (num_steps,)
. This should have event_shape(hidden_dim,)
.  observation_matrix (Tensor) – A linear transformation from hidden
to observed state. This should have shape broadcastable to
self.batch_shape + (num_steps, hidden_dim, obs_dim)
.  observation_dist – A observation noise distribution. This should have
batch_shape broadcastable to
self.batch_shape + (num_steps,)
. This should have event_shape(obs_dim,)
.

arg_constraints
= {}¶

has_rsample
= True¶

support
¶
 initial_dist – A distribution over initial states. This should have
batch_shape broadcastable to
LKJCorrCholesky¶

class
LKJCorrCholesky
(d, eta, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Generates cholesky factors of correlation matrices using an LKJ prior.
The expected use is to combine it with a vector of variances and pass it to the scale_tril parameter of a multivariate distribution such as MultivariateNormal.
E.g., if theta is a (positive) vector of covariances with the same dimensionality as this distribution, and Omega is sampled from this distribution, scale_tril=torch.mm(torch.diag(sqrt(theta)), Omega)
Note that the event_shape of this distribution is [d, d]
Note
When using this distribution with HMC/NUTS, it is important to use a step_size such as 1e4. If not, you are likely to experience LAPACK errors regarding positivedefiniteness.
For example usage, refer to pyro/examples/lkj.py.
Parameters:  d (int) – Dimensionality of the matrix
 eta (torch.Tensor) – A single positive number parameterizing the distribution.

arg_constraints
= {'eta': GreaterThan(lower_bound=0.0)}¶

has_rsample
= False¶

support
= CorrCholesky()¶
MaskedMixture¶

class
MaskedMixture
(mask, component0, component1, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
A masked deterministic mixture of two distributions.
This is useful when the mask is sampled from another distribution, possibly correlated across the batch. Often the mask can be marginalized out via enumeration.
Example:
change_point = pyro.sample("change_point", dist.Categorical(torch.ones(len(data) + 1)), infer={'enumerate': 'parallel'}) mask = torch.arange(len(data), dtype=torch.long) >= changepoint with pyro.plate("data", len(data)): pyro.sample("obs", MaskedMixture(mask, dist1, dist2), obs=data)
Parameters:  mask (torch.Tensor) – A byte tensor toggling between
component0
andcomponent1
.  component0 (pyro.distributions.TorchDistribution) – a distribution
for batch elements
mask == 0
.  component1 (pyro.distributions.TorchDistribution) – a distribution
for batch elements
mask == 1
.

arg_constraints
= {}¶

has_rsample
¶

support
¶
 mask (torch.Tensor) – A byte tensor toggling between
MixtureOfDiagNormals¶

class
MixtureOfDiagNormals
(locs, coord_scale, component_logits)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Mixture of Normal distributions with arbitrary means and arbitrary diagonal covariance matrices.
That is, this distribution is a mixture with K components, where each component distribution is a Ddimensional Normal distribution with a Ddimensional mean parameter and a Ddimensional diagonal covariance matrix. The K different component means are gathered into the K x D dimensional parameter locs and the K different scale parameters are gathered into the K x D dimensional parameter coord_scale. The mixture weights are controlled by a Kdimensional vector of softmax logits, component_logits. This distribution implements pathwise derivatives for samples from the distribution.
See reference [1] for details on the implementations of the pathwise derivative. Please consider citing this reference if you use the pathwise derivative in your research. Note that this distribution does not support dimension D = 1.
[1] Pathwise Derivatives for Multivariate Distributions, Martin Jankowiak & Theofanis Karaletsos. arXiv:1806.01856
Parameters:  locs (torch.Tensor) – K x D mean matrix
 coord_scale (torch.Tensor) – K x D scale matrix
 component_logits (torch.Tensor) – Kdimensional vector of softmax logits

arg_constraints
= {'component_logits': Real(), 'coord_scale': GreaterThan(lower_bound=0.0), 'locs': Real()}¶

has_rsample
= True¶
MultivariateStudentT¶

class
MultivariateStudentT
(df, loc, scale_tril, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Creates a multivariate Student’s tdistribution parameterized by degree of freedom
df
, meanloc
and scalescale_tril
.Parameters: 
arg_constraints
= {'df': GreaterThan(lower_bound=0.0), 'loc': RealVector(), 'scale_tril': LowerCholesky()}¶

has_rsample
= True¶

mean
¶

support
= RealVector()¶

variance
¶

OMTMultivariateNormal¶

class
OMTMultivariateNormal
(loc, scale_tril)[source]¶ Bases:
pyro.distributions.torch.MultivariateNormal
Multivariate normal (Gaussian) distribution with OMT gradients w.r.t. both parameters. Note the gradient computation w.r.t. the Cholesky factor has cost O(D^3), although the resulting gradient variance is generally expected to be lower.
A distribution over vectors in which all the elements have a joint Gaussian density.
Parameters:  loc (torch.Tensor) – Mean.
 scale_tril (torch.Tensor) – Cholesky of Covariance matrix.

arg_constraints
= {'loc': Real(), 'scale_tril': LowerTriangular()}¶
RelaxedBernoulliStraightThrough¶

class
RelaxedBernoulliStraightThrough
(temperature, probs=None, logits=None, validate_args=None)[source]¶ Bases:
pyro.distributions.torch.RelaxedBernoulli
An implementation of
RelaxedBernoulli
with a straightthrough gradient estimator.This distribution has the following properties:
 The samples returned by the
rsample()
method are discrete/quantized.  The
log_prob()
method returns the log probability of the relaxed/unquantized sample using the GumbelSoftmax distribution.  In the backward pass the gradient of the sample with respect to the parameters of the distribution uses the relaxed/unquantized sample.
References:
 [1] The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables,
 Chris J. Maddison, Andriy Mnih, Yee Whye Teh
 [2] Categorical Reparameterization with GumbelSoftmax,
 Eric Jang, Shixiang Gu, Ben Poole
 The samples returned by the
RelaxedOneHotCategoricalStraightThrough¶

class
RelaxedOneHotCategoricalStraightThrough
(temperature, probs=None, logits=None, validate_args=None)[source]¶ Bases:
pyro.distributions.torch.RelaxedOneHotCategorical
An implementation of
RelaxedOneHotCategorical
with a straightthrough gradient estimator.This distribution has the following properties:
 The samples returned by the
rsample()
method are discrete/quantized.  The
log_prob()
method returns the log probability of the relaxed/unquantized sample using the GumbelSoftmax distribution.  In the backward pass the gradient of the sample with respect to the parameters of the distribution uses the relaxed/unquantized sample.
References:
 [1] The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables,
 Chris J. Maddison, Andriy Mnih, Yee Whye Teh
 [2] Categorical Reparameterization with GumbelSoftmax,
 Eric Jang, Shixiang Gu, Ben Poole
 The samples returned by the
Rejector¶

class
Rejector
(propose, log_prob_accept, log_scale)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Rejection sampled distribution given an acceptance rate function.
Parameters:  propose (Distribution) – A proposal distribution that samples batched
proposals via
propose()
.rsample()
supports asample_shape
arg only ifpropose()
supports asample_shape
arg.  log_prob_accept (callable) – A callable that inputs a batch of proposals and returns a batch of log acceptance probabilities.
 log_scale – Total log probability of acceptance.

has_rsample
= True¶
 propose (Distribution) – A proposal distribution that samples batched
proposals via
SpanningTree¶

class
SpanningTree
(edge_logits, sampler_options=None, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Distribution over spanning trees on a fixed number
V
of vertices.A tree is represented as
torch.LongTensor
edges
of shape(V1,2)
satisfying the following properties: The edges constitute a tree, i.e. are connected and cycle free.
 Each edge
(v1,v2) = edges[e]
is sorted, i.e.v1 < v2
.  The entire tensor is sorted in colexicographic order.
Use
validate_edges()
to verify edges are correctly formed.The
edge_logits
tensor has one entry for each of theV*(V1)//2
edges in the complete graph onV
vertices, where edges are each sorted and the edge order is colexicographic:(0,1), (0,2), (1,2), (0,3), (1,3), (2,3), (0,4), (1,4), (2,4), ...
This ordering corresponds to the sizeindependent pairing function:
k = v1 + v2 * (v2  1) // 2
where
k
is the rank of the edge(v1,v2)
in the complete graph. To convert a matrix of edge logits to the linear representation used here:assert my_matrix.shape == (V, V) i, j = make_complete_graph(V) edge_logits = my_matrix[i, j]
Parameters:  edge_logits (torch.Tensor) – A tensor of length
V*(V1)//2
containing logits (aka negative energies) of all edges in the complete graph onV
vertices. See above comment for edge ordering.  sampler_options (dict) – An optional dict of sampler options including:
mcmc_steps
defaulting to a single MCMC step (which is pretty good);initial_edges
defaulting to a cheap approximate sample;backend
one of “python” or “cpp”, defaulting to “python”.

arg_constraints
= {'edge_logits': Real()}¶

enumerate_support
(expand=True)[source]¶ This is implemented for trees with up to 6 vertices (and 5 edges).

has_enumerate_support
= True¶

sample
(sample_shape=torch.Size([]))[source]¶ This sampler is implemented using MCMC run for a small number of steps after being initialized by a cheap approximate sampler. This sampler is approximate and cubic time. This is faster than the classic AldousBroder sampler [1,2], especially for graphs with large mixing time. Recent research [3,4] proposes samplers that run in submatrixmultiply time but are more complex to implement.
References
 [1] Generating random spanning trees
 Andrei Broder (1989)
 [2] The Random Walk Construction of Uniform Spanning Trees and Uniform Labelled Trees,
 David J. Aldous (1990)
 [3] Sampling Random Spanning Trees Faster than Matrix Multiplication,
 David Durfee, Rasmus Kyng, John Peebles, Anup B. Rao, Sushant Sachdeva (2017) https://arxiv.org/abs/1611.07451
 [4] An almostlinear time algorithm for uniform random spanning tree generation,
 Aaron Schild (2017) https://arxiv.org/abs/1711.06455

support
= IntegerGreaterThan(lower_bound=0)¶

validate_edges
(edges)[source]¶ Validates a batch of
edges
tensors, as returned bysample()
orenumerate_support()
or as input tolog_prob()
.Parameters: edges (torch.LongTensor) – A batch of edges. Raises: ValueError Returns: None
Stable¶

class
Stable
(stability, skew, scale=1.0, loc=0.0, coords='S0', validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Levy \(\alpha\)stable distribution. See [1] for a review.
This uses Nolan’s parametrization [2] of the
loc
parameter, which is required for continuity and differentiability. This corresponds to the notation \(S^0_\alpha(\beta,\sigma,\mu_0)\) of [1], where \(\alpha\) = stability, \(\beta\) = skew, \(\sigma\) = scale, and \(\mu_0\) = loc. To instead use the S parameterization as in scipy, passcoords="S"
, but BEWARE this is discontinuous atstability=1
and has poor geometry for inference.This implements a reparametrized sampler
rsample()
, but does not implementlog_prob()
. Inference can be performed using either likelihoodfree algorithms such asEnergyDistance
, or reparameterization via thereparam()
handler with one of the reparameterizersLatentStableReparam
,SymmetricStableReparam
, orStableReparam
e.g.:with poutine.reparam(config={"x": StableReparam()}): pyro.sample("x", Stable(stability, skew, scale, loc))
 [1] S. Borak, W. Hardle, R. Weron (2005).
 Stable distributions. https://edoc.huberlin.de/bitstream/handle/18452/4526/8.pdf
 [2] J.P. Nolan (1997).
 Numerical calculation of stable densities and distribution functions.
 [3] Rafal Weron (1996).
 On the ChambersMallowsStuck Method for Simulating Skewed Stable Random Variables.
 [4] J.P. Nolan (2017).
 Stable Distributions: Models for Heavy Tailed Data. http://fs2.american.edu/jpnolan/www/stable/chap1.pdf
Parameters:  stability (Tensor) – Levy stability parameter \(\alpha\in(0,2]\) .
 skew (Tensor) – Skewness \(\beta\in[1,1]\) .
 scale (Tensor) – Scale \(\sigma > 0\) . Defaults to 1.
 loc (Tensor) – Location \(\mu_0\) when using Nolan’s S0 parametrization [2], or \(\mu\) when using the S parameterization. Defaults to 0.
 coords (str) – Either “S0” (default) to use Nolan’s continuous S0 parametrization, or “S” to use the discontinuous parameterization.

arg_constraints
= {'loc': Real(), 'scale': GreaterThan(lower_bound=0.0), 'skew': Interval(lower_bound=1, upper_bound=1), 'stability': Interval(lower_bound=0, upper_bound=2)}¶

has_rsample
= True¶

support
= Real()¶
Unit¶

class
Unit
(log_factor, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Trivial nonnormalized distribution representing the unit type.
The unit type has a single value with no data, i.e.
value.numel() == 0
.This is used for
pyro.factor()
statements.
arg_constraints
= {'log_factor': Real()}¶

support
= Real()¶

VonMises¶

class
VonMises
(loc, concentration, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
A circular von Mises distribution.
This implementation uses polar coordinates. The
loc
andvalue
args can be any real number (to facilitate unconstrained optimization), but are interpreted as angles modulo 2 pi.See
VonMises3D
for a 3D cartesian coordinate cousin of this distribution.Parameters:  loc (torch.Tensor) – an angle in radians.
 concentration (torch.Tensor) – concentration parameter

arg_constraints
= {'concentration': GreaterThan(lower_bound=0.0), 'loc': Real()}¶

has_rsample
= False¶

mean
¶ The provided mean is the circular one.

sample
(sample_shape=torch.Size([]))[source]¶ The sampling algorithm for the von Mises distribution is based on the following paper: Best, D. J., and Nicholas I. Fisher. “Efficient simulation of the von Mises distribution.” Applied Statistics (1979): 152157.

support
= Real()¶

variance
¶ The provided variance is the circular one.
VonMises3D¶

class
VonMises3D
(concentration, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Spherical von Mises distribution.
This implementation combines the direction parameter and concentration parameter into a single combined parameter that contains both direction and magnitude. The
value
arg is represented in cartesian coordinates: it must be a normalized 3vector that lies on the 2sphere.See
VonMises
for a 2D polar coordinate cousin of this distribution.Currently only
log_prob()
is implemented.Parameters: concentration (torch.Tensor) – A combined locationandconcentration vector. The direction of this vector is the location, and its magnitude is the concentration. 
arg_constraints
= {'concentration': Real()}¶

support
= Real()¶

ZeroInflatedPoisson¶

class
ZeroInflatedPoisson
(gate, rate, validate_args=None)[source]¶ Bases:
pyro.distributions.zero_inflated.ZeroInflatedDistribution
A Zero Inflated Poisson distribution.
Parameters:  gate (torch.Tensor) – probability of extra zeros.
 rate (torch.Tensor) – rate of poisson distribution.

support
= IntegerGreaterThan(lower_bound=0)¶
ZeroInflatedNegativeBinomial¶

class
ZeroInflatedNegativeBinomial
(gate, total_count, probs=None, logits=None, validate_args=None)[source]¶ Bases:
pyro.distributions.zero_inflated.ZeroInflatedDistribution
A Zero Inflated Negative Binomial distribution.
Parameters:  gate (torch.Tensor) – probability of extra zeros.
 total_count (float or torch.Tensor) – nonnegative number of negative Bernoulli trials.
 probs (torch.Tensor) – Event probabilities of success in the half open interval [0, 1).
 logits (torch.Tensor) – Event logodds for probabilities of success.

support
= IntegerGreaterThan(lower_bound=0)¶
ZeroInflatedDistribution¶

class
ZeroInflatedDistribution
(gate, base_dist, validate_args=None)[source]¶ Bases:
pyro.distributions.torch_distribution.TorchDistribution
Base class for a Zero Inflated distribution.
Parameters:  gate (torch.Tensor) – probability of extra zeros given via a Bernoulli distribution.
 base_dist (TorchDistribution) – the base distribution.

arg_constraints
= {'gate': Interval(lower_bound=0.0, upper_bound=1.0)}¶
Transforms¶
ConditionalTransform¶
CorrLCholeskyTransform¶

class
CorrLCholeskyTransform
(cache_size=0)[source]¶ Bases:
torch.distributions.transforms.Transform
Transforms a vector into the cholesky factor of a correlation matrix.
The input should have shape [batch_shape] + [d * (d1)/2]. The output will have shape [batch_shape] + [d, d].
Reference:
[1] Cholesky Factors of Correlation Matrices, Stan Reference Manual v2.18, Section 10.12

bijective
= True¶

codomain
= CorrCholesky()¶

domain
= Real()¶

event_dim
= 1¶

sign
= 1¶

ELUTransform¶
LeakyReLUTransform¶
LowerCholeskyAffine¶

class
LowerCholeskyAffine
(loc, scale_tril)[source]¶ Bases:
torch.distributions.transforms.Transform
A bijection of the form \(\mathbf{y} = \mathbf{L} \mathbf{x} + \mathbf{r}\) where mathbf{L} is a lower triangular matrix and mathbf{r} is a vector.
Parameters:  loc (torch.tensor) – the fixed Ddimensional vector to shift the input by.
 scale_tril (torch.tensor) – the D x D lower triangular matrix used in the transformation.

bijective
= True¶

codomain
= RealVector()¶

event_dim
= 1¶

log_abs_det_jacobian
(x, y)[source]¶ Calculates the elementwise determinant of the log Jacobian, i.e. log(abs(dy/dx)).

volume_preserving
= False¶
Permute¶

class
Permute
(permutation)[source]¶ Bases:
torch.distributions.transforms.Transform
A bijection that reorders the input dimensions, that is, multiplies the input by a permutation matrix. This is useful in between
AffineAutoregressive
transforms to increase the flexibility of the resulting distribution and stabilize learning. Whilst not being an autoregressive transform, the log absolute determinate of the Jacobian is easily calculable as 0. Note that reordering the input dimension between two layers ofAffineAutoregressive
is not equivalent to reordering the dimension inside the MADE networks that those IAFs use; using aPermute
transform results in a distribution with more flexibility.Example usage:
>>> from pyro.nn import AutoRegressiveNN >>> from pyro.distributions.transforms import AffineAutoregressive, Permute >>> base_dist = dist.Normal(torch.zeros(10), torch.ones(10)) >>> iaf1 = AffineAutoregressive(AutoRegressiveNN(10, [40])) >>> ff = Permute(torch.randperm(10, dtype=torch.long)) >>> iaf2 = AffineAutoregressive(AutoRegressiveNN(10, [40])) >>> flow_dist = dist.TransformedDistribution(base_dist, [iaf1, ff, iaf2]) >>> flow_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
Parameters: permutation (torch.LongTensor) – a permutation ordering that is applied to the inputs. 
bijective
= True¶

codomain
= Real()¶

event_dim
= 1¶

log_abs_det_jacobian
(x, y)[source]¶ Calculates the elementwise determinant of the log Jacobian, i.e. log(abs([dy_0/dx_0, …, dy_{N1}/dx_{N1}])). Note that this type of transform is not autoregressive, so the log Jacobian is not the sum of the previous expression. However, it turns out it’s always 0 (since the determinant is 1 or +1), and so returning a vector of zeros works.

volume_preserving
= True¶

TanhTransform¶
DiscreteCosineTransform¶

class
DiscreteCosineTransform
(dim=1, cache_size=0)[source]¶ Bases:
torch.distributions.transforms.Transform
Discrete Cosine Transform of typeII.
This uses
dct()
andidct()
to compute orthonormal DCT and inverse DCT transforms. The jacobian is 1.Parameters: dim (int) – Dimension along which to transform. Must be negative. 
bijective
= True¶

codomain
= Real()¶

domain
= Real()¶

TransformModules¶
AffineAutoregressive¶

class
AffineAutoregressive
(autoregressive_nn, log_scale_min_clip=5.0, log_scale_max_clip=3.0, sigmoid_bias=2.0, stable=False)[source]¶ Bases:
pyro.distributions.torch_transform.TransformModule
An implementation of the bijective transform of Inverse Autoregressive Flow (IAF), using by default Eq (10) from Kingma Et Al., 2016,
\(\mathbf{y} = \mu_t + \sigma_t\odot\mathbf{x}\)where \(\mathbf{x}\) are the inputs, \(\mathbf{y}\) are the outputs, \(\mu_t,\sigma_t\) are calculated from an autoregressive network on \(\mathbf{x}\), and \(\sigma_t>0\).
If the stable keyword argument is set to True then the transformation used is,
\(\mathbf{y} = \sigma_t\odot\mathbf{x} + (1\sigma_t)\odot\mu_t\)where \(\sigma_t\) is restricted to \((0,1)\). This variant of IAF is claimed by the authors to be more numerically stable than one using Eq (10), although in practice it leads to a restriction on the distributions that can be represented, presumably since the input is restricted to rescaling by a number on \((0,1)\).
Together with
TransformedDistribution
this provides a way to create richer variational approximations.Example usage:
>>> from pyro.nn import AutoRegressiveNN >>> base_dist = dist.Normal(torch.zeros(10), torch.ones(10)) >>> transform = AffineAutoregressive(AutoRegressiveNN(10, [40])) >>> pyro.module("my_transform", transform) # doctest: +SKIP >>> flow_dist = dist.TransformedDistribution(base_dist, [transform]) >>> flow_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
The inverse of the Bijector is required when, e.g., scoring the log density of a sample with
TransformedDistribution
. This implementation caches the inverse of the Bijector when its forward operation is called, e.g., when sampling fromTransformedDistribution
. However, if the cached value isn’t available, either because it was overwritten during sampling a new value or an arbitary value is being scored, it will calculate it manually. Note that this is an operation that scales as O(D) where D is the input dimension, and so should be avoided for large dimensional uses. So in general, it is cheap to sample from IAF and score a value that was sampled by IAF, but expensive to score an arbitrary value.Parameters:  autoregressive_nn (nn.Module) – an autoregressive neural network whose forward call returns a realvalued mean and logitscale as a tuple
 log_scale_min_clip (float) – The minimum value for clipping the log(scale) from the autoregressive NN
 log_scale_max_clip (float) – The maximum value for clipping the log(scale) from the autoregressive NN
 sigmoid_bias (float) – A term to add the logit of the input when using the stable tranform.
 stable (bool) – When true, uses the alternative “stable” version of the transform (see above).
References:
1. Improving Variational Inference with Inverse Autoregressive Flow [arXiv:1606.04934] Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, Max Welling
2. Variational Inference with Normalizing Flows [arXiv:1505.05770] Danilo Jimenez Rezende, Shakir Mohamed
3. MADE: Masked Autoencoder for Distribution Estimation [arXiv:1502.03509] Mathieu Germain, Karol Gregor, Iain Murray, Hugo Larochelle

autoregressive
= True¶

bijective
= True¶

codomain
= Real()¶

domain
= Real()¶

event_dim
= 1¶
AffineCoupling¶

class
AffineCoupling
(split_dim, hypernet, log_scale_min_clip=5.0, log_scale_max_clip=3.0)[source]¶ Bases:
pyro.distributions.torch_transform.TransformModule
An implementation of the affine coupling layer of RealNVP (Dinh et al., 2017) that uses the bijective transform,
\(\mathbf{y}_{1:d} = \mathbf{x}_{1:d}\) \(\mathbf{y}_{(d+1):D} = \mu + \sigma\odot\mathbf{x}_{(d+1):D}\)where \(\mathbf{x}\) are the inputs, \(\mathbf{y}\) are the outputs, e.g. \(\mathbf{x}_{1:d}\) represents the first \(d\) elements of the inputs, and \(\mu,\sigma\) are shift and translation parameters calculated as the output of a function inputting only \(\mathbf{x}_{1:d}\).
That is, the first \(d\) components remain unchanged, and the subsequent \(Dd\) are shifted and translated by a function of the previous components.
Together with
TransformedDistribution
this provides a way to create richer variational approximations.Example usage:
>>> from pyro.nn import DenseNN >>> input_dim = 10 >>> split_dim = 6 >>> base_dist = dist.Normal(torch.zeros(input_dim), torch.ones(input_dim)) >>> hypernet = DenseNN(split_dim, [10*input_dim], [input_dimsplit_dim, input_dimsplit_dim]) >>> transform = AffineCoupling(split_dim, hypernet) >>> pyro.module("my_transform", transform) # doctest: +SKIP >>> flow_dist = dist.TransformedDistribution(base_dist, [transform]) >>> flow_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
The inverse of the Bijector is required when, e.g., scoring the log density of a sample with
TransformedDistribution
. This implementation caches the inverse of the Bijector when its forward operation is called, e.g., when sampling fromTransformedDistribution
. However, if the cached value isn’t available, either because it was overwritten during sampling a new value or an arbitary value is being scored, it will calculate it manually.This is an operation that scales as O(1), i.e. constant in the input dimension. So in general, it is cheap to sample and score (an arbitrary value) from
AffineCoupling
.Parameters:  split_dim (int) – Zeroindexed dimension \(d\) upon which to perform input/output split for transformation.
 hypernet (callable) – an autoregressive neural network whose forward call returns a realvalued mean and logitscale as a tuple. The input should have final dimension split_dim and the output final dimension input_dimsplit_dim for each member of the tuple.
 log_scale_min_clip (float) – The minimum value for clipping the log(scale) from the autoregressive NN
 log_scale_max_clip (float) – The maximum value for clipping the log(scale) from the autoregressive NN
References:
Laurent Dinh, Jascha SohlDickstein, and Samy Bengio. Density estimation using Real NVP. ICLR 2017.

bijective
= True¶

codomain
= Real()¶

domain
= Real()¶

event_dim
= 1¶
BatchNorm¶

class
BatchNorm
(input_dim, momentum=0.1, epsilon=1e05)[source]¶ Bases:
pyro.distributions.torch_transform.TransformModule
A type of batch normalization that can be used to stabilize training in normalizing flows. The inverse operation is defined as
\(x = (y  \hat{\mu}) \oslash \sqrt{\hat{\sigma^2}} \otimes \gamma + \beta\)that is, the standard batch norm equation, where \(x\) is the input, \(y\) is the output, \(\gamma,\beta\) are learnable parameters, and \(\hat{\mu}\)/\(\hat{\sigma^2}\) are smoothed running averages of the sample mean and variance, respectively. The constraint \(\gamma>0\) is enforced to ease calculation of the logdetJacobian term.
This is an elementwise transform, and when applied to a vector, learns two parameters (\(\gamma,\beta\)) for each dimension of the input.
When the module is set to training mode, the moving averages of the sample mean and variance are updated every time the inverse operator is called, e.g., when a normalizing flow scores a minibatch with the log_prob method.
Also, when the module is set to training mode, the sample mean and variance on the current minibatch are used in place of the smoothed averages, \(\hat{\mu}\) and \(\hat{\sigma^2}\), for the inverse operator. For this reason it is not the case that \(x=g(g^{1}(x))\) during training, i.e., that the inverse operation is the inverse of the forward one.
Example usage:
>>> from pyro.nn import AutoRegressiveNN >>> from pyro.distributions.transforms import AffineAutoregressive >>> base_dist = dist.Normal(torch.zeros(10), torch.ones(10)) >>> iafs = [AffineAutoregressive(AutoRegressiveNN(10, [40])) for _ in range(2)] >>> bn = BatchNorm(10) >>> flow_dist = dist.TransformedDistribution(base_dist, [iafs[0], bn, iafs[1]]) >>> flow_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
Parameters: References:
[1] Sergey Ioffe and Christian Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. In International Conference on Machine Learning, 2015. https://arxiv.org/abs/1502.03167
[2] Laurent Dinh, Jascha SohlDickstein, and Samy Bengio. Density Estimation using Real NVP. In International Conference on Learning Representations, 2017. https://arxiv.org/abs/1605.08803
[3] George Papamakarios, Theo Pavlakou, and Iain Murray. Masked Autoregressive Flow for Density Estimation. In Neural Information Processing Systems, 2017. https://arxiv.org/abs/1705.07057

bijective
= True¶

codomain
= Real()¶

constrained_gamma
¶

domain
= Real()¶

event_dim
= 0¶

BlockAutoregressive¶

class
BlockAutoregressive
(input_dim, hidden_factors=[8, 8], activation='tanh', residual=None)[source]¶ Bases:
pyro.distributions.torch_transform.TransformModule
An implementation of Block Neural Autoregressive Flow (blockNAF) (De Cao et al., 2019) bijective transform. BlockNAF uses a similar transformation to deep dense NAF, building the autoregressive NN into the structure of the transform, in a sense.
Together with
TransformedDistribution
this provides a way to create richer variational approximations.Example usage:
>>> base_dist = dist.Normal(torch.zeros(10), torch.ones(10)) >>> naf = BlockAutoregressive(input_dim=10) >>> pyro.module("my_naf", naf) # doctest: +SKIP >>> naf_dist = dist.TransformedDistribution(base_dist, [naf]) >>> naf_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
The inverse operation is not implemented. This would require numerical inversion, e.g., using a root finding method  a possibility for a future implementation.
Parameters:  input_dim (int) – The dimensionality of the input and output variables.
 hidden_factors (list) – Hidden layer i has hidden_factors[i] hidden units per input dimension. This corresponds to both \(a\) and \(b\) in De Cao et al. (2019). The elements of hidden_factors must be integers.
 activation (string) – Activation function to use. One of ‘ELU’, ‘LeakyReLU’, ‘sigmoid’, or ‘tanh’.
 residual (string) – Type of residual connections to use. Choices are “None”, “normal” for \(\mathbf{y}+f(\mathbf{y})\), and “gated” for \(\alpha\mathbf{y} + (1  \alpha\mathbf{y})\) for learnable parameter \(\alpha\).
References:
Block Neural Autoregressive Flow [arXiv:1904.04676] Nicola De Cao, Ivan Titov, Wilker Aziz

autoregressive
= True¶

bijective
= True¶

codomain
= Real()¶

domain
= Real()¶

event_dim
= 1¶
ConditionalPlanar¶

class
ConditionalPlanar
(nn)[source]¶ Bases:
pyro.distributions.conditional.ConditionalTransformModule
A conditional ‘planar’ bijective transform using the equation,
\(\mathbf{y} = \mathbf{x} + \mathbf{u}\tanh(\mathbf{w}^T\mathbf{z}+b)\)where \(\mathbf{x}\) are the inputs with dimension \(D\), \(\mathbf{y}\) are the outputs, and the pseudoparameters \(b\in\mathbb{R}\), \(\mathbf{u}\in\mathbb{R}^D\), and \(\mathbf{w}\in\mathbb{R}^D\) are the output of a function, e.g. a NN, with input \(z\in\mathbb{R}^{M}\) representing the context variable to condition on. For this to be an invertible transformation, the condition \(\mathbf{w}^T\mathbf{u}>1\) is enforced.
Together with
ConditionalTransformedDistribution
this provides a way to create richer variational approximations.Example usage:
>>> from pyro.nn.dense_nn import DenseNN >>> input_dim = 10 >>> context_dim = 5 >>> batch_size = 3 >>> base_dist = dist.Normal(torch.zeros(input_dim), torch.ones(input_dim)) >>> hypernet = DenseNN(context_dim, [50, 50], param_dims=[1, input_dim, input_dim]) >>> transform = ConditionalPlanar(hypernet) >>> z = torch.rand(batch_size, context_dim) >>> flow_dist = dist.ConditionalTransformedDistribution(base_dist, [transform]).condition(z) >>> flow_dist.sample(sample_shape=torch.Size([batch_size])) # doctest: +SKIP
The inverse of this transform does not possess an analytical solution and is left unimplemented. However, the inverse is cached when the forward operation is called during sampling, and so samples drawn using the planar transform can be scored.
Parameters: nn (callable) – a function inputting the context variable and outputting a triplet of realvalued parameters of dimensions \((1, D, D)\). References: Variational Inference with Normalizing Flows [arXiv:1505.05770] Danilo Jimenez Rezende, Shakir Mohamed

bijective
= True¶

codomain
= Real()¶

condition
(context)[source]¶ See
pyro.distributions.conditional.ConditionalTransformModule.condition()

domain
= Real()¶

event_dim
= 1¶

ConditionalTransformModule¶

class
ConditionalTransformModule
(*args, **kwargs)[source]¶ Bases:
pyro.distributions.conditional.ConditionalTransform
,torch.nn.modules.module.Module
Conditional transforms with learnable parameters such as normalizing flows should inherit from this class rather than
ConditionalTransform
so they are also a subclass ofModule
and inherit all the useful methods of that class.
Householder¶

class
Householder
(input_dim, count_transforms=1)[source]¶ Bases:
pyro.distributions.torch_transform.TransformModule
Represents multiple applications of the Householder bijective transformation. A single Householder transformation takes the form,
\(\mathbf{y} = (I  2*\frac{\mathbf{u}\mathbf{u}^T}{\mathbf{u}^2})\mathbf{x}\)where \(\mathbf{x}\) are the inputs, \(\mathbf{y}\) are the outputs, and the learnable parameters are \(\mathbf{u}\in\mathbb{R}^D\) for input dimension \(D\).
The transformation represents the reflection of \(\mathbf{x}\) through the plane passing through the origin with normal \(\mathbf{u}\).
\(D\) applications of this transformation are able to transform standard i.i.d. standard Gaussian noise into a Gaussian variable with an arbitrary covariance matrix. With \(K<D\) transformations, one is able to approximate a fullrank Gaussian distribution using a linear transformation of rank \(K\).
Together with
TransformedDistribution
this provides a way to create richer variational approximations.Example usage:
>>> base_dist = dist.Normal(torch.zeros(10), torch.ones(10)) >>> transform = Householder(10, count_transforms=5) >>> pyro.module("my_transform", p) # doctest: +SKIP >>> flow_dist = dist.TransformedDistribution(base_dist, [transform]) >>> flow_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
Parameters: References:
Improving Variational AutoEncoders using Householder Flow, [arXiv:1611.09630] Tomczak, J. M., & Welling, M.

bijective
= True¶

codomain
= Real()¶

domain
= Real()¶

event_dim
= 1¶

log_abs_det_jacobian
(x, y)[source]¶ Calculates the elementwise determinant of the log jacobian. Householder flow is measure preserving, so \(\log(detJ) = 0\)

volume_preserving
= True¶

NeuralAutoregressive¶

class
NeuralAutoregressive
(autoregressive_nn, hidden_units=16, activation='sigmoid')[source]¶ Bases:
pyro.distributions.torch_transform.TransformModule
An implementation of the deep Neural Autoregressive Flow (NAF) bijective transform of the “IAF flavour” that can be used for sampling and scoring samples drawn from it (but not arbitrary ones).
Example usage:
>>> from pyro.nn import AutoRegressiveNN >>> base_dist = dist.Normal(torch.zeros(10), torch.ones(10)) >>> arn = AutoRegressiveNN(10, [40], param_dims=[16]*3) >>> transform = NeuralAutoregressive(arn, hidden_units=16) >>> pyro.module("my_transform", transform) # doctest: +SKIP >>> flow_dist = dist.TransformedDistribution(base_dist, [transform]) >>> flow_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
The inverse operation is not implemented. This would require numerical inversion, e.g., using a root finding method  a possibility for a future implementation.
Parameters:  autoregressive_nn (nn.Module) – an autoregressive neural network whose forward call returns a tuple of three realvalued tensors, whose last dimension is the input dimension, and whose penultimate dimension is equal to hidden_units.
 hidden_units (int) – the number of hidden units to use in the NAF transformation (see Eq (8) in reference)
 activation (string) – Activation function to use. One of ‘ELU’, ‘LeakyReLU’, ‘sigmoid’, or ‘tanh’.
Reference:
Neural Autoregressive Flows [arXiv:1804.00779] ChinWei Huang, David Krueger, Alexandre Lacoste, Aaron Courville

bijective
= True¶

codomain
= Real()¶

domain
= Real()¶

event_dim
= 1¶
Planar¶

class
Planar
(input_dim)[source]¶ Bases:
pyro.distributions.transforms.planar.ConditionedPlanar
,pyro.distributions.torch_transform.TransformModule
A ‘planar’ bijective transform with equation,
\(\mathbf{y} = \mathbf{x} + \mathbf{u}\tanh(\mathbf{w}^T\mathbf{z}+b)\)where \(\mathbf{x}\) are the inputs, \(\mathbf{y}\) are the outputs, and the learnable parameters are \(b\in\mathbb{R}\), \(\mathbf{u}\in\mathbb{R}^D\), \(\mathbf{w}\in\mathbb{R}^D\) for input dimension \(D\). For this to be an invertible transformation, the condition \(\mathbf{w}^T\mathbf{u}>1\) is enforced.
Together with
TransformedDistribution
this provides a way to create richer variational approximations.Example usage:
>>> base_dist = dist.Normal(torch.zeros(10), torch.ones(10)) >>> transform = Planar(10) >>> pyro.module("my_transform", transform) # doctest: +SKIP >>> flow_dist = dist.TransformedDistribution(base_dist, [transform]) >>> flow_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
The inverse of this transform does not possess an analytical solution and is left unimplemented. However, the inverse is cached when the forward operation is called during sampling, and so samples drawn using the planar transform can be scored.
Parameters: input_dim (int) – the dimension of the input (and output) variable. References:
Variational Inference with Normalizing Flows [arXiv:1505.05770] Danilo Jimenez Rezende, Shakir Mohamed

bijective
= True¶

codomain
= Real()¶

domain
= Real()¶

event_dim
= 1¶

Polynomial¶

class
Polynomial
(autoregressive_nn, input_dim, count_degree, count_sum)[source]¶ Bases:
pyro.distributions.torch_transform.TransformModule
An autoregressive bijective transform as described in Jaini et al. (2019) applying following equation elementwise,
\(y_n = c_n + \int^{x_n}_0\sum^K_{k=1}\left(\sum^R_{r=0}a^{(n)}_{r,k}u^r\right)du\)where \(x_n\) is the \(n\) is the \(n\), \(\left\{a^{(n)}_{r,k}\in\mathbb{R}\right\}\) are learnable parameters that are the output of an autoregressive NN inputting \(x_{\prec n}={x_1,x_2,\ldots,x_{n1}}\).
Together with
TransformedDistribution
this provides a way to create richer variational approximations.Example usage:
>>> from pyro.nn import AutoRegressiveNN >>> input_dim = 10 >>> count_degree = 4 >>> count_sum = 3 >>> base_dist = dist.Normal(torch.zeros(input_dim), torch.ones(input_dim)) >>> arn = AutoRegressiveNN(input_dim, [input_dim*10], param_dims=[(count_degree + 1)*count_sum]) >>> transform = Polynomial(arn, input_dim=input_dim, count_degree=count_degree, count_sum=count_sum) >>> pyro.module("my_transform", transform) # doctest: +SKIP >>> flow_dist = dist.TransformedDistribution(base_dist, [transform]) >>> flow_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
The inverse of this transform does not possess an analytical solution and is left unimplemented. However, the inverse is cached when the forward operation is called during sampling, and so samples drawn using a polynomial transform can be scored.
Parameters:  autoregressive_nn (nn.Module) – an autoregressive neural network whose forward call returns a tensor of realvalued numbers of size (batch_size, (count_degree+1)*count_sum, input_dim)
 count_degree (int) – The degree of the polynomial to use for each elementwise transformation.
 count_sum (int) – The number of polynomials to sum in each elementwise transformation.
References:
Sumofsquares polynomial flow. [arXiv:1905.02325] Priyank Jaini, Kira A. Shelby, Yaoliang Yu

autoregressive
= True¶

bijective
= True¶

codomain
= Real()¶

domain
= Real()¶

event_dim
= 1¶
Radial¶

class
Radial
(input_dim)[source]¶ Bases:
pyro.distributions.torch_transform.TransformModule
A ‘radial’ bijective transform using the equation,
\(\mathbf{y} = \mathbf{x} + \beta h(\alpha,r)(\mathbf{x}  \mathbf{x}_0)\)where \(\mathbf{x}\) are the inputs, \(\mathbf{y}\) are the outputs, and the learnable parameters are \(\alpha\in\mathbb{R}^+\), \(\beta\in\mathbb{R}\), \(\mathbf{x}_0\in\mathbb{R}^D\), for input dimension \(D\), \(r=\mathbf{x}\mathbf{x}_0_2\), \(h(\alpha,r)=1/(\alpha+r)\). For this to be an invertible transformation, the condition \(\beta>\alpha\) is enforced.
Together with
TransformedDistribution
this provides a way to create richer variational approximations.Example usage:
>>> base_dist = dist.Normal(torch.zeros(10), torch.ones(10)) >>> transform = Radial(10) >>> pyro.module("my_transform", transform) # doctest: +SKIP >>> flow_dist = dist.TransformedDistribution(base_dist, [transform]) >>> flow_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
The inverse of this transform does not possess an analytical solution and is left unimplemented. However, the inverse is cached when the forward operation is called during sampling, and so samples drawn using the radial transform can be scored.
Parameters: input_dim (int) – the dimension of the input (and output) variable. References:
Variational Inference with Normalizing Flows [arXiv:1505.05770] Danilo Jimenez Rezende, Shakir Mohamed

bijective
= True¶

codomain
= Real()¶

domain
= Real()¶

event_dim
= 1¶

Sylvester¶

class
Sylvester
(input_dim, count_transforms=1)[source]¶ Bases:
pyro.distributions.transforms.householder.Householder
An implementation of the Sylvester bijective transform of the Householder variety (Van den Berg Et Al., 2018),
\(\mathbf{y} = \mathbf{x} + QR\tanh(SQ^T\mathbf{x}+\mathbf{b})\)where \(\mathbf{x}\) are the inputs, \(\mathbf{y}\) are the outputs, \(R,S\sim D\times D\) are upper triangular matrices for input dimension \(D\), \(Q\sim D\times D\) is an orthogonal matrix, and \(\mathbf{b}\sim D\) is learnable bias term.
The Sylvester transform is a generalization of
Planar
. In the Householder type of the Sylvester transform, the orthogonality of \(Q\) is enforced by representing it as the product of Householder transformations.Together with
TransformedDistribution
it provides a way to create richer variational approximations.Example usage:
>>> base_dist = dist.Normal(torch.zeros(10), torch.ones(10)) >>> transform = Sylvester(10, count_transforms=4) >>> pyro.module("my_transform", transform) # doctest: +SKIP >>> flow_dist = dist.TransformedDistribution(base_dist, [transform]) >>> flow_dist.sample() # doctest: +SKIP tensor([0.4071, 0.5030, 0.7924, 0.2366, 0.2387, 0.1417, 0.0868, 0.1389, 0.4629, 0.0986])
The inverse of this transform does not possess an analytical solution and is left unimplemented. However, the inverse is cached when the forward operation is called during sampling, and so samples drawn using the Sylvester transform can be scored.
References:
Rianne van den Berg, Leonard Hasenclever, Jakub M. Tomczak, Max Welling. Sylvester Normalizing Flows for Variational Inference. In proceedings of The 34th Conference on Uncertainty in Artificial Intelligence (UAI 2018).

bijective
= True¶

codomain
= Real()¶

domain
= Real()¶

event_dim
= 1¶

TransformModule¶

class
TransformModule
(*args, **kwargs)[source]¶ Bases:
torch.distributions.transforms.Transform
,torch.nn.modules.module.Module
Transforms with learnable parameters such as normalizing flows should inherit from this class rather than Transform so they are also a subclass of nn.Module and inherit all the useful methods of that class.
Transform Factories¶
Each Transform
and TransformModule
includes a corresponding helper function in lower case that inputs, at minimum, the input dimensions of the transform, and possibly additional arguments to customize the transform in an intuitive way. The purpose of these helper functions is to hide from the user whether or not the transform requires the construction of a hypernet, and if so, the input and output dimensions of that hypernet.
affine_autoregressive¶

affine_autoregressive
(input_dim, hidden_dims=None, **kwargs)[source]¶ A helper function to create an
AffineAutoregressive
object that takes care of constructing an autoregressive network with the correct input/output dimensions.Parameters:  input_dim (int) – Dimension of input variable
 hidden_dims (list[int]) – The desired hidden dimensions of the autoregressive network. Defaults to using [3*input_dim + 1]
 log_scale_min_clip (float) – The minimum value for clipping the log(scale) from the autoregressive NN
 log_scale_max_clip (float) – The maximum value for clipping the log(scale) from the autoregressive NN
 sigmoid_bias (float) – A term to add the logit of the input when using the stable tranform.
 stable (bool) – When true, uses the alternative “stable” version of the transform (see above).
affine_coupling¶

affine_coupling
(input_dim, hidden_dims=None, split_dim=None, **kwargs)[source]¶ A helper function to create an
AffineCoupling
object that takes care of constructing a dense network with the correct input/output dimensions.Parameters:  input_dim (int) – Dimension of input variable
 hidden_dims (list[int]) – The desired hidden dimensions of the dense network. Defaults to using [10*input_dim]
 split_dim (int) – The dimension to split the input on for the coupling transform. Defaults to using input_dim // 2
 log_scale_min_clip (float) – The minimum value for clipping the log(scale) from the autoregressive NN
 log_scale_max_clip (float) – The maximum value for clipping the log(scale) from the autoregressive NN
batchnorm¶
block_autoregressive¶

block_autoregressive
(input_dim, **kwargs)[source]¶ A helper function to create a
BlockAutoregressive
object for consistency with other helpers.Parameters:  input_dim (int) – Dimension of input variable
 hidden_factors (list) – Hidden layer i has hidden_factors[i] hidden units per input dimension. This corresponds to both \(a\) and \(b\) in De Cao et al. (2019). The elements of hidden_factors must be integers.
 activation (string) – Activation function to use. One of ‘ELU’, ‘LeakyReLU’, ‘sigmoid’, or ‘tanh’.
 residual (string) – Type of residual connections to use. Choices are “None”, “normal” for \(\mathbf{y}+f(\mathbf{y})\), and “gated” for \(\alpha\mathbf{y} + (1  \alpha\mathbf{y})\) for learnable parameter \(\alpha\).
conditional_planar¶

conditional_planar
(input_dim, context_dim, hidden_dims=None)[source]¶ A helper function to create a
ConditionalPlanar
object that takes care of constructing a dense network with the correct input/output dimensions.Parameters:
elu¶
householder¶

householder
(input_dim, count_transforms=None)[source]¶ A helper function to create a
Householder
object for consistency with other helpers.Parameters:
leaky_relu¶

leaky_relu
()[source]¶ A helper function to create a
LeakyReLUTransform
object for consistency with other helpers.
neural_autoregressive¶

neural_autoregressive
(input_dim, hidden_dims=None, activation='sigmoid', width=16)[source]¶ A helper function to create a
NeuralAutoregressive
object that takes care of constructing an autoregressive network with the correct input/output dimensions.Parameters:  input_dim (int) – Dimension of input variable
 hidden_dims (list[int]) – The desired hidden dimensions of the autoregressive network. Defaults to using [3*input_dim + 1]
 activation (string) – Activation function to use. One of ‘ELU’, ‘LeakyReLU’, ‘sigmoid’, or ‘tanh’.
 width (int) – The width of the “multilayer perceptron” in the transform (see paper). Defaults to 16
permute¶

permute
(input_dim, permutation=None)[source]¶ A helper function to create a
Permute
object for consistency with other helpers.Parameters:  input_dim (int) – Dimension of input variable
 permutation (torch.LongTensor) – Torch tensor of integer indices representing permutation. Defaults to a random permutation.
planar¶
polynomial¶

polynomial
(input_dim, hidden_dims=None)[source]¶ A helper function to create a
Polynomial
object that takes care of constructing an autoregressive network with the correct input/output dimensions.Parameters:  input_dim (int) – Dimension of input variable
 hidden_dims – The desired hidden dimensions of of the autoregressive network. Defaults to using [input_dim * 10]
radial¶
sylvester¶
tanh¶

tanh
()[source]¶ A helper function to create a
TanhTransform
object for consistency with other helpers.
Parameters¶
Parameters in Pyro are basically thin wrappers around PyTorch Tensors that carry unique names. As such Parameters are the primary stateful objects in Pyro. Users typically interact with parameters via the Pyro primitive pyro.param. Parameters play a central role in stochastic variational inference, where they are used to represent point estimates for the parameters in parameterized families of models and guides.
ParamStore¶

class
ParamStoreDict
[source]¶ Bases:
object
Global store for parameters in Pyro. This is basically a keyvalue store. The typical user interacts with the ParamStore primarily through the primitive pyro.param.
See Intro Part II for further discussion and SVI Part I for some examples.
Some things to bear in mind when using parameters in Pyro:
 parameters must be assigned unique names
 the init_tensor argument to pyro.param is only used the first time that a given (named) parameter is registered with Pyro.
 for this reason, a user may need to use the clear() method if working in a REPL in order to get the desired behavior. this method can also be invoked with pyro.clear_param_store().
 the internal name of a parameter within a PyTorch nn.Module that has been registered with Pyro is prepended with the Pyro name of the module. so nothing prevents the user from having two different modules each of which contains a parameter named weight. by contrast, a user can only have one toplevel parameter named weight (outside of any module).
 parameters can be saved and loaded from disk using save and load.

setdefault
(name, init_constrained_value, constraint=Real())[source]¶ Retrieve a constrained parameter value from the if it exists, otherwise set the initial value. Note that this is a little fancier than
dict.setdefault()
.If the parameter already exists,
init_constrained_tensor
will be ignored. To avoid expensive creation ofinit_constrained_tensor
you can wrap it in alambda
that will only be evaluated if the parameter does not already exist:param_store.get("foo", lambda: (0.001 * torch.randn(1000, 1000)).exp(), constraint=constraints.positive)
Parameters:  name (str) – parameter name
 init_constrained_value (torch.Tensor or callable returning a torch.Tensor) – initial constrained value
 constraint (torch.distributions.constraints.Constraint) – torch constraint object
Returns: constrained parameter value
Return type:

named_parameters
()[source]¶ Returns an iterator over
(name, unconstrained_value)
tuples for each parameter in the ParamStore.

get_param
(name, init_tensor=None, constraint=Real(), event_dim=None)[source]¶ Get parameter from its name. If it does not yet exist in the ParamStore, it will be created and stored. The Pyro primitive pyro.param dispatches to this method.
Parameters:  name (str) – parameter name
 init_tensor (torch.Tensor) – initial tensor
 constraint (torch.distributions.constraints.Constraint) – torch constraint
 event_dim (int) – (ignored)
Returns: parameter
Return type:

match
(name)[source]¶ Get all parameters that match regex. The parameter must exist.
Parameters: name (str) – regular expression Returns: dict with key param name and value torch Tensor

param_name
(p)[source]¶ Get parameter name from parameter
Parameters: p – parameter Returns: parameter name

load
(filename, map_location=None)[source]¶ Loads parameters from disk
Note
If using
pyro.module()
on parameters loaded from disk, be sure to set theupdate_module_params
flag:pyro.get_param_store().load('saved_params.save') pyro.module('module', nn, update_module_params=True)
Parameters:  filename (str) – file name to load from
 map_location (function, torch.device, string or a dict) – specifies how to remap storage locations
Neural Networks¶
The module pyro.nn provides implementations of neural network modules that are useful in the context of deep probabilistic programming.
Pyro Modules¶
Pyro includes an experimental class PyroModule
, a
subclass of torch.nn.Module
, whose attributes can be modified by Pyro
effects. To create a poutineaware attribute, use either the
PyroParam
struct or the PyroSample
struct:
my_module = PyroModule()
my_module.x = PyroParam(torch.tensor(1.), constraint=constraints.positive)
my_module.y = PyroSample(dist.Normal(0, 1))

class
PyroParam
¶ Bases:
tuple
Structure to declare a Pyromanaged learnable parameter of a
PyroModule
.
constraint
¶ Alias for field number 1

event_dim
¶ Alias for field number 2

init_value
¶ Alias for field number 0


class
PyroSample
¶ Bases:
tuple
Structure to declare a Pyromanaged random parameter of a
PyroModule
.
prior
¶ Alias for field number 0


class
PyroModule
(name='')[source]¶ Bases:
torch.nn.modules.module.Module
EXPERIMENTAL Subclass of
torch.nn.Module
whose attributes can be modified by Pyro effects. Attributes can be set using helpersPyroParam
andPyroSample
, and methods can be decorated bypyro_method()
.Parameters
To create a Pyromanaged parameter attribute, set that attribute using either
torch.nn.Parameter
(for unconstrained parameters) orPyroParam
(for constrained parameters). Reading that attribute will then trigger apyro.param()
statement. For example:# Create Pyromanaged parameter attributes. my_module = PyroModule() my_module.loc = nn.Parameter(torch.tensor(0.)) my_module.scale = PyroParam(torch.tensor(1.), constraint=constraints.positive) # Read the attributes. loc = my_module.loc # Triggers a pyro.param statement. scale = my_module.scale # Triggers another pyro.param statement.
Note that, unlike normal
torch.nn.Module
s,PyroModule
s should not be registered withpyro.module()
statements.PyroModule
s can contain otherPyroModule
s and normaltorch.nn.Module
s. Accessing a normaltorch.nn.Module
attribute of aPyroModule
triggers apyro.module()
statement. If multiplePyroModule
s appear in a single Pyro model or guide, they should be included in a single rootPyroModule
for that model.PyroModule
s synchronize data with the param store at eachsetattr
,getattr
, anddelattr
event, based on the nested name of an attribute: Setting
mod.x = x_init
tries to readx
from the param store. If a value is found in the param store, that value is copied intomod
andx_init
is ignored; otherwisex_init
is copied into bothmod
and the param store.  Reading
mod.x
tries to readx
from the param store. If a value is found in the param store, that value is copied intomod
; otherwisemod
’s value is copied into the param store. Finallymod
and the param store agree on a single value to return.  Deleting
del mod.x
removes a value from bothmod
and the param store.
Note two
PyroModule
of the same name will both synchronize with the global param store and thus contain the same data. When creating aPyroModule
, then deleting it, then creating another with the same name, the latter will be populated with the former’s data from the param store. To avoid this persistence, eitherpyro.clear_param_store()
or callclear()
before deleting aPyroModule
.PyroModule
s can be saved and loaded either directly usingtorch.save()
/torch.load()
or indirectly using the param store’ssave()
/load()
. Note thattorch.load()
will be overridden by any values in the param store, so it is safest topyro.clear_param_store()
before loading.Samples
To create a Pyromanaged random attribute, set that attribute using the
PyroSample
helper, specifying a prior distribution. Reading that attribute will then trigger apyro.sample()
statement. For example:# Create Pyromanaged random attributes. my_module.x = PyroSample(dist.Normal(0, 1)) my_module.y = PyroSample(lambda self: dist.Normal(self.loc, self.scale)) # Sample the attributes. x = my_module.x # Triggers a pyro.sample statement. y = my_module.y # Triggers one pyro.sample + two pyro.param statements.
Sampling is cached within each invocation of
.__call__()
or method decorated bypyro_method()
. Because sample statements can appear only once in a Pyro trace, you should ensure that traced access to sample attributes is wrapped in a single invocation of.__call__()
or method decorated bypyro_method()
.To make an existing module probabilistic, you can create a subclass and overwrite some parameters with
PyroSample
s:class RandomLinear(nn.Linear, PyroModule): # used as a mixin def __init__(self, in_features, out_features): super().__init__(in_features, out_features) self.weight = PyroSample( lambda self: dist.Normal(0, 1) .expand([self.out_features, self.in_features]) .to_event(2))
Mixin classes
PyroModule
can be used as a mixin class, and supports simple syntax for dynamically creating mixins, for example the following are equivalent:# Version 1. create a named mixin class class PyroLinear(nn.Linear, PyroModule): pass m.linear = PyroLinear(m, n) # Version 2. create a dynamic mixin class m.linear = PyroModule[nn.Linear](m, n)
This notation can be used recursively to create Bayesian modules, e.g.:
model = PyroModule[nn.Sequential]( PyroModule[nn.Linear](28 * 28, 100), PyroModule[nn.Sigmoid](), PyroModule[nn.Linear](100, 100), PyroModule[nn.Sigmoid](), PyroModule[nn.Linear](100, 10), ) assert isinstance(model, nn.Sequential) assert isinstance(model, PyroModule) # Now we can be Bayesian about weights in the first layer. model[0].weight = PyroSample( prior=dist.Normal(0, 1).expand([28 * 28, 100]).to_event(2)) guide = AutoDiagonalNormal(model)
Note that
PyroModule[...]
does not recursively mix inPyroModule
to submodules of the inputModule
; hence we needed to wrap each submodule of thenn.Sequential
above.Parameters: name (str) – Optional name for a root PyroModule. This is ignored in subPyroModules of another PyroModule.  Setting

pyro_method
(fn)[source]¶ Decorator for toplevel methods of a
PyroModule
to enable pyro effects and cachepyro.sample
statements.This should be applied to all public methods that read Pyromanaged attributes, but is not needed for
.forward()
.

clear
(mod)[source]¶ Removes data from both a
PyroModule
and the param store.Parameters: mod (PyroModule) – A module to clear.

to_pyro_module_
(m, recurse=True)[source]¶ Converts an ordinary
torch.nn.Module
instance to aPyroModule
inplace.This is useful for adding Pyro effects to thirdparty modules: no thirdparty code needs to be modified. For example:
model = nn.Sequential( nn.Linear(28 * 28, 100), nn.Sigmoid(), nn.Linear(100, 100), nn.Sigmoid(), nn.Linear(100, 10), ) to_pyro_module_(model) assert isinstance(model, PyroModule[nn.Sequential]) assert isinstance(model[0], PyroModule[nn.Linear]) # Now we can attempt to be fully Bayesian: for m in model.modules(): for name, value in list(m.named_parameters(recurse=False)): setattr(m, name, PyroSample(prior=dist.Normal(0, 1) .expand(value.shape) .to_event(value.dim()))) guide = AutoDiagonalNormal(model)
Parameters:  m (torch.nn.Module) – A module instance.
 recurse (bool) – Whether to convert submodules to
PyroModules
.
AutoRegressiveNN¶

class
AutoRegressiveNN
(input_dim, hidden_dims, param_dims=[1, 1], permutation=None, skip_connections=False, nonlinearity=ReLU())[source]¶ Bases:
pyro.nn.auto_reg_nn.ConditionalAutoRegressiveNN
An implementation of a MADElike autoregressive neural network.
Example usage:
>>> x = torch.randn(100, 10) >>> arn = AutoRegressiveNN(10, [50], param_dims=[1]) >>> p = arn(x) # 1 parameters of size (100, 10) >>> arn = AutoRegressiveNN(10, [50], param_dims=[1, 1]) >>> m, s = arn(x) # 2 parameters of size (100, 10) >>> arn = AutoRegressiveNN(10, [50], param_dims=[1, 5, 3]) >>> a, b, c = arn(x) # 3 parameters of sizes, (100, 1, 10), (100, 5, 10), (100, 3, 10)
Parameters:  input_dim (int) – the dimensionality of the input variable
 hidden_dims (list[int]) – the dimensionality of the hidden units per layer
 param_dims (list[int]) – shape the output into parameters of dimension (p_n, input_dim) for p_n in param_dims when p_n > 1 and dimension (input_dim) when p_n == 1. The default is [1, 1], i.e. output two parameters of dimension (input_dim), which is useful for inverse autoregressive flow.
 permutation (torch.LongTensor) – an optional permutation that is applied to the inputs and controls the order of the autoregressive factorization. in particular for the identity permutation the autoregressive structure is such that the Jacobian is upper triangular. By default this is chosen at random.
 skip_connections (bool) – Whether to add skip connections from the input to the output.
 nonlinearity (torch.nn.module) – The nonlinearity to use in the feedforward network such as torch.nn.ReLU(). Note that no nonlinearity is applied to the final network output, so the output is an unbounded real number.
Reference:
MADE: Masked Autoencoder for Distribution Estimation [arXiv:1502.03509] Mathieu Germain, Karol Gregor, Iain Murray, Hugo Larochelle
ConditionalAutoRegressiveNN¶

class
ConditionalAutoRegressiveNN
(input_dim, context_dim, hidden_dims, param_dims=[1, 1], permutation=None, skip_connections=False, nonlinearity=ReLU())[source]¶ Bases:
torch.nn.modules.module.Module
An implementation of a MADElike autoregressive neural network that can input an additional context variable. (See Reference [2] Section 3.3 for an explanation of how the conditional MADE architecture works.)
Example usage:
>>> x = torch.randn(100, 10) >>> y = torch.randn(100, 5) >>> arn = ConditionalAutoRegressiveNN(10, 5, [50], param_dims=[1]) >>> p = arn(x, context=y) # 1 parameters of size (100, 10) >>> arn = ConditionalAutoRegressiveNN(10, 5, [50], param_dims=[1, 1]) >>> m, s = arn(x, context=y) # 2 parameters of size (100, 10) >>> arn = ConditionalAutoRegressiveNN(10, 5, [50], param_dims=[1, 5, 3]) >>> a, b, c = arn(x, context=y) # 3 parameters of sizes, (100, 1, 10), (100, 5, 10), (100, 3, 10)
Parameters:  input_dim (int) – the dimensionality of the input variable
 context_dim (int) – the dimensionality of the context variable
 hidden_dims (list[int]) – the dimensionality of the hidden units per layer
 param_dims (list[int]) – shape the output into parameters of dimension (p_n, input_dim) for p_n in param_dims when p_n > 1 and dimension (input_dim) when p_n == 1. The default is [1, 1], i.e. output two parameters of dimension (input_dim), which is useful for inverse autoregressive flow.
 permutation (torch.LongTensor) – an optional permutation that is applied to the inputs and controls the order of the autoregressive factorization. in particular for the identity permutation the autoregressive structure is such that the Jacobian is upper triangular. By default this is chosen at random.
 skip_connections (bool) – Whether to add skip connections from the input to the output.
 nonlinearity (torch.nn.module) – The nonlinearity to use in the feedforward network such as torch.nn.ReLU(). Note that no nonlinearity is applied to the final network output, so the output is an unbounded real number.
Reference:
1. MADE: Masked Autoencoder for Distribution Estimation [arXiv:1502.03509] Mathieu Germain, Karol Gregor, Iain Murray, Hugo Larochelle
2. Inference Networks for Sequential Monte Carlo in Graphical Models [arXiv:1602.06701] Brooks Paige, Frank Wood
Optimization¶
The module pyro.optim provides support for optimization in Pyro. In particular it provides PyroOptim, which is used to wrap PyTorch optimizers and manage optimizers for dynamically generated parameters (see the tutorial SVI Part I for a discussion). Any custom optimization algorithms are also to be found here.
Pyro Optimizers¶

class
PyroOptim
(optim_constructor, optim_args, clip_args=None)[source]¶ Bases:
object
A wrapper for torch.optim.Optimizer objects that helps with managing dynamically generated parameters.
Parameters:  optim_constructor – a torch.optim.Optimizer
 optim_args – a dictionary of learning arguments for the optimizer or a callable that returns such dictionaries
 clip_args – a dictionary of clip_norm and/or clip_value args or a callable that returns such dictionaries

__call__
(params, *args, **kwargs)[source]¶ Parameters: params (an iterable of strings) – a list of parameters Do an optimization step for each param in params. If a given param has never been seen before, initialize an optimizer for it.

get_state
()[source]¶ Get state associated with all the optimizers in the form of a dictionary with keyvalue pairs (parameter name, optim state dicts)

set_state
(state_dict)[source]¶ Set the state associated with all the optimizers using the state obtained from a previous call to get_state()

AdagradRMSProp
(optim_args)[source]¶ Wraps
pyro.optim.adagrad_rmsprop.AdagradRMSProp
withPyroOptim
.

ClippedAdam
(optim_args)[source]¶ Wraps
pyro.optim.clipped_adam.ClippedAdam
withPyroOptim
.

class
PyroLRScheduler
(scheduler_constructor, optim_args, clip_args=None)[source]¶ Bases:
pyro.optim.optim.PyroOptim
A wrapper for
lr_scheduler
objects that adjusts learning rates for dynamically generated parameters.Parameters:  scheduler_constructor – a
lr_scheduler
 optim_args – a dictionary of learning arguments for the optimizer or a callable that returns such dictionaries. must contain the key ‘optimizer’ with pytorch optimizer value
 clip_args – a dictionary of clip_norm and/or clip_value args or a callable that returns such dictionaries.
Example:
optimizer = torch.optim.SGD scheduler = pyro.optim.ExponentialLR({'optimizer': optimizer, 'optim_args': {'lr': 0.01}, 'gamma': 0.1}) svi = SVI(model, guide, scheduler, loss=TraceGraph_ELBO()) for i in range(epochs): for minibatch in DataLoader(dataset, batch_size): svi.step(minibatch) scheduler.step(epoch=i)
 scheduler_constructor – a

class
AdagradRMSProp
(params, eta=1.0, delta=1e16, t=0.1)[source]¶ Bases:
torch.optim.optimizer.Optimizer
Implements a mashup of the Adagrad algorithm and RMSProp. For the precise update equation see equations 10 and 11 in reference [1].
References: [1] ‘Automatic Differentiation Variational Inference’, Alp Kucukelbir, Dustin Tran, Rajesh Ranganath, Andrew Gelman, David M. Blei URL: https://arxiv.org/abs/1603.00788 [2] ‘Lecture 6.5 RmsProp: Divide the gradient by a running average of its recent magnitude’, Tieleman, T. and Hinton, G., COURSERA: Neural Networks for Machine Learning. [3] ‘Adaptive subgradient methods for online learning and stochastic optimization’, Duchi, John, Hazan, E and Singer, Y.
Arguments:
Parameters:  params – iterable of parameters to optimize or dicts defining parameter groups
 eta (float) – sets the step size scale (optional; default: 1.0)
 t (float) – t, optional): momentum parameter (optional; default: 0.1)
 delta (float) – modulates the exponent that controls how the step size scales (optional: default: 1e16)

class
ClippedAdam
(params, lr=0.001, betas=(0.9, 0.999), eps=1e08, weight_decay=0, clip_norm=10.0, lrd=1.0)[source]¶ Bases:
torch.optim.optimizer.Optimizer
Parameters:  params – iterable of parameters to optimize or dicts defining parameter groups
 lr – learning rate (default: 1e3)
 betas (Tuple) – coefficients used for computing running averages of gradient and its square (default: (0.9, 0.999))
 eps – term added to the denominator to improve numerical stability (default: 1e8)
 weight_decay – weight decay (L2 penalty) (default: 0)
 clip_norm – magnitude of norm to which gradients are clipped (default: 10.0)
 lrd – rate at which learning rate decays (default: 1.0)
Small modification to the Adam algorithm implemented in torch.optim.Adam to include gradient clipping and learning rate decay.
Reference
A Method for Stochastic Optimization, Diederik P. Kingma, Jimmy Ba https://arxiv.org/abs/1412.6980
PyTorch Optimizers¶

Adadelta
(optim_args, clip_args=None)¶ Wraps
torch.optim.Adadelta
withPyroOptim
.

Adagrad
(optim_args, clip_args=None)¶ Wraps
torch.optim.Adagrad
withPyroOptim
.

Adam
(optim_args, clip_args=None)¶ Wraps
torch.optim.Adam
withPyroOptim
.

AdamW
(optim_args, clip_args=None)¶ Wraps
torch.optim.AdamW
withPyroOptim
.

SparseAdam
(optim_args, clip_args=None)¶ Wraps
torch.optim.SparseAdam
withPyroOptim
.

Adamax
(optim_args, clip_args=None)¶ Wraps
torch.optim.Adamax
withPyroOptim
.

ASGD
(optim_args, clip_args=None)¶ Wraps
torch.optim.ASGD
withPyroOptim
.

SGD
(optim_args, clip_args=None)¶ Wraps
torch.optim.SGD
withPyroOptim
.

Rprop
(optim_args, clip_args=None)¶ Wraps
torch.optim.Rprop
withPyroOptim
.

RMSprop
(optim_args, clip_args=None)¶ Wraps
torch.optim.RMSprop
withPyroOptim
.

LambdaLR
(optim_args, clip_args=None)¶ Wraps
torch.optim.LambdaLR
withPyroLRScheduler
.

MultiplicativeLR
(optim_args, clip_args=None)¶ Wraps
torch.optim.MultiplicativeLR
withPyroLRScheduler
.

StepLR
(optim_args, clip_args=None)¶ Wraps
torch.optim.StepLR
withPyroLRScheduler
.

MultiStepLR
(optim_args, clip_args=None)¶ Wraps
torch.optim.MultiStepLR
withPyroLRScheduler
.

ExponentialLR
(optim_args, clip_args=None)¶ Wraps
torch.optim.ExponentialLR
withPyroLRScheduler
.

CosineAnnealingLR
(optim_args, clip_args=None)¶ Wraps
torch.optim.CosineAnnealingLR
withPyroLRScheduler
.

ReduceLROnPlateau
(optim_args, clip_args=None)¶ Wraps
torch.optim.ReduceLROnPlateau
withPyroLRScheduler
.

CyclicLR
(optim_args, clip_args=None)¶ Wraps
torch.optim.CyclicLR
withPyroLRScheduler
.

CosineAnnealingWarmRestarts
(optim_args, clip_args=None)¶ Wraps
torch.optim.CosineAnnealingWarmRestarts
withPyroLRScheduler
.

OneCycleLR
(optim_args, clip_args=None)¶ Wraps
torch.optim.OneCycleLR
withPyroLRScheduler
.
HigherOrder Optimizers¶

class
MultiOptimizer
[source]¶ Bases:
object
Base class of optimizers that make use of higherorder derivatives.
Higherorder optimizers generally use
torch.autograd.grad()
rather thantorch.Tensor.backward()
, and therefore require a different interface from usual Pyro and PyTorch optimizers. In this interface, thestep()
method inputs aloss
tensor to be differentiated, and backpropagation is triggered one or more times inside the optimizer.Derived classes must implement
step()
to compute derivatives and update parameters inplace.Example:
tr = poutine.trace(model).get_trace(*args, **kwargs) loss = tr.log_prob_sum() params = {name: site['value'].unconstrained() for name, site in tr.nodes.items() if site['type'] == 'param'} optim.step(loss, params)

step
(loss, params)[source]¶ Performs an inplace optimization step on parameters given a differentiable
loss
tensor.Note that this detaches the updated tensors.
Parameters:  loss (torch.Tensor) – A differentiable tensor to be minimized. Some optimizers require this to be differentiable multiple times.
 params (dict) – A dictionary mapping param name to unconstrained value as stored in the param store.

get_step
(loss, params)[source]¶ Computes an optimization step of parameters given a differentiable
loss
tensor, returning the updated values.Note that this preserves derivatives on the updated tensors.
Parameters:  loss (torch.Tensor) – A differentiable tensor to be minimized. Some optimizers require this to be differentiable multiple times.
 params (dict) – A dictionary mapping param name to unconstrained value as stored in the param store.
Returns: A dictionary mapping param name to updated unconstrained value.
Return type:


class
PyroMultiOptimizer
(optim)[source]¶ Bases:
pyro.optim.multi.MultiOptimizer
Facade to wrap
PyroOptim
objects in aMultiOptimizer
interface.

class
TorchMultiOptimizer
(optim_constructor, optim_args)[source]¶ Bases:
pyro.optim.multi.PyroMultiOptimizer
Facade to wrap
Optimizer
objects in aMultiOptimizer
interface.

class
MixedMultiOptimizer
(parts)[source]¶ Bases:
pyro.optim.multi.MultiOptimizer
Container class to combine different
MultiOptimizer
instances for different parameters.Parameters: parts (list) – A list of (names, optim)
pairs, where eachnames
is a list of parameter names, and eachoptim
is aMultiOptimizer
orPyroOptim
object to be used for the named parameters. Together thenames
should partition up all desired parameters to optimize.Raises: ValueError – if any name is optimized by multiple optimizers.

class
Newton
(trust_radii={})[source]¶ Bases:
pyro.optim.multi.MultiOptimizer
Implementation of
MultiOptimizer
that performs a Newton update on batched lowdimensional variables, optionally regularizing via a perparametertrust_radius
. Seenewton_step()
for details.The result of
get_step()
will be differentiable, however the updated values fromstep()
will be detached.Parameters: trust_radii (dict) – a dict mapping parameter name to radius of trust region. Missing names will use unregularized Newton update, equivalent to infinite trust radius.
Poutine (Effect handlers)¶
Beneath the builtin inference algorithms, Pyro has a library of composable effect handlers for creating new inference algorithms and working with probabilistic programs. Pyro’s inference algorithms are all built by applying these handlers to stochastic functions.
Handlers¶
Poutine is a library of composable effect handlers for recording and modifying the behavior of Pyro programs. These lowerlevel ingredients simplify the implementation of new inference algorithms and behavior.
Handlers can be used as higherorder functions, decorators, or context managers to modify the behavior of functions or blocks of code:
For example, consider the following Pyro program:
>>> def model(x):
... s = pyro.param("s", torch.tensor(0.5))
... z = pyro.sample("z", dist.Normal(x, s))
... return z ** 2
We can mark sample sites as observed using condition
,
which returns a callable with the same input and output signatures as model
:
>>> conditioned_model = poutine.condition(model, data={"z": 1.0})
We can also use handlers as decorators:
>>> @pyro.condition(data={"z": 1.0})
... def model(x):
... s = pyro.param("s", torch.tensor(0.5))
... z = pyro.sample("z", dist.Normal(x, s))
... return z ** 2
Or as context managers:
>>> with pyro.condition(data={"z": 1.0}):
... s = pyro.param("s", torch.tensor(0.5))
... z = pyro.sample("z", dist.Normal(0., s))
... y = z ** 2
Handlers compose freely:
>>> conditioned_model = poutine.condition(model, data={"z": 1.0})
>>> traced_model = poutine.trace(conditioned_model)
Many inference algorithms or algorithmic components can be implemented in just a few lines of code:
guide_tr = poutine.trace(guide).get_trace(...)
model_tr = poutine.trace(poutine.replay(conditioned_model, trace=guide_tr)).get_trace(...)
monte_carlo_elbo = model_tr.log_prob_sum()  guide_tr.log_prob_sum()

block
(fn=None, *args, **kwargs)¶ Convenient wrapper of
BlockMessenger
This handler selectively hides Pyro primitive sites from the outside world. Default behavior: block everything.
A site is hidden if at least one of the following holds:
hide_fn(msg) is True
or(not expose_fn(msg)) is True
msg["name"] in hide
msg["type"] in hide_types
msg["name"] not in expose and msg["type"] not in expose_types
hide
,hide_types
, andexpose_types
are allNone
For example, suppose the stochastic function fn has two sample sites “a” and “b”. Then any effect outside of
BlockMessenger(fn, hide=["a"])
will not be applied to site “a” and will only see site “b”:>>> def fn(): ... a = pyro.sample("a", dist.Normal(0., 1.)) ... return pyro.sample("b", dist.Normal(a, 1.)) >>> fn_inner = pyro.poutine.trace(fn) >>> fn_outer = pyro.poutine.trace(pyro.poutine.block(fn_inner, hide=["a"])) >>> trace_inner = fn_inner.get_trace() >>> trace_outer = fn_outer.get_trace() >>> "a" in trace_inner True >>> "a" in trace_outer False >>> "b" in trace_inner True >>> "b" in trace_outer True
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 hide_fn – function that takes a site and returns True to hide the site or False/None to expose it. If specified, all other parameters are ignored. Only specify one of hide_fn or expose_fn, not both.
 expose_fn – function that takes a site and returns True to expose the site or False/None to hide it. If specified, all other parameters are ignored. Only specify one of hide_fn or expose_fn, not both.
 hide_all (bool) – hide all sites
 expose_all (bool) – expose all sites normally
 hide (list) – list of site names to hide
 expose (list) – list of site names to be exposed while all others hidden
 hide_types (list) – list of site types to be hidden
 expose_types (lits) – list of site types to be exposed while all others hidden
Returns: stochastic function decorated with a
BlockMessenger

broadcast
(fn=None, *args, **kwargs)¶ Convenient wrapper of
BroadcastMessenger
Automatically broadcasts the batch shape of the stochastic function at a sample site when inside a single or nested plate context. The existing batch_shape must be broadcastable with the size of the
plate
contexts installed in the cond_indep_stack.Notice how model_automatic_broadcast below automates expanding of distribution batch shapes. This makes it easy to modularize a Pyro model as the subcomponents are agnostic of the wrapping
plate
contexts.>>> def model_broadcast_by_hand(): ... with IndepMessenger("batch", 100, dim=2): ... with IndepMessenger("components", 3, dim=1): ... sample = pyro.sample("sample", dist.Bernoulli(torch.ones(3) * 0.5) ... .expand_by(100)) ... assert sample.shape == torch.Size((100, 3)) ... return sample
>>> @poutine.broadcast ... def model_automatic_broadcast(): ... with IndepMessenger("batch", 100, dim=2): ... with IndepMessenger("components", 3, dim=1): ... sample = pyro.sample("sample", dist.Bernoulli(torch.tensor(0.5))) ... assert sample.shape == torch.Size((100, 3)) ... return sample

condition
(fn=None, *args, **kwargs)¶ Convenient wrapper of
ConditionMessenger
Given a stochastic function with some sample statements and a dictionary of observations at names, change the sample statements at those names into observes with those values.
Consider the following Pyro program:
>>> def model(x): ... s = pyro.param("s", torch.tensor(0.5)) ... z = pyro.sample("z", dist.Normal(x, s)) ... return z ** 2
To observe a value for site z, we can write
>>> conditioned_model = pyro.poutine.condition(model, data={"z": torch.tensor(1.)})
This is equivalent to adding obs=value as a keyword argument to pyro.sample(“z”, …) in model.
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 data – a dict or a
Trace
Returns: stochastic function decorated with a
ConditionMessenger

do
(fn=None, *args, **kwargs)¶ Convenient wrapper of
DoMessenger
Given a stochastic function with some sample statements and a dictionary of values at names, set the return values of those sites equal to the values as if they were hardcoded to those values and introduce fresh sample sites with the same names whose values do not propagate.
Composes freely with
condition()
to represent counterfactual distributions over potential outcomes. See Single World Intervention Graphs [1] for additional details and theory.Consider the following Pyro program:
>>> def model(x): ... s = pyro.param("s", torch.tensor(0.5)) ... z = pyro.sample("z", dist.Normal(x, s)) ... return z ** 2
To intervene with a value for site z, we can write
>>> intervened_model = pyro.poutine.do(model, data={"z": torch.tensor(1.)})
This is equivalent to replacing z = pyro.sample(“z”, …) with z = torch.tensor(1.) and introducing a fresh sample site pyro.sample(“z”, …) whose value is not used elsewhere.
References
 [1] Single World Intervention Graphs: A Primer,
 Thomas Richardson, James Robins
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 data – a
dict
mapping sample site names to interventions
Returns: stochastic function decorated with a
DoMessenger

enum
(fn=None, *args, **kwargs)¶ Convenient wrapper of
EnumMessenger
Enumerates in parallel over discrete sample sites marked
infer={"enumerate": "parallel"}
.Parameters: first_available_dim (int) – The first tensor dimension (counting from the right) that is available for parallel enumeration. This dimension and all dimensions left may be used internally by Pyro. This should be a negative integer or None.

escape
(fn=None, *args, **kwargs)¶ Convenient wrapper of
EscapeMessenger
Messenger that does a nonlocal exit by raising a util.NonlocalExit exception

infer_config
(fn=None, *args, **kwargs)¶ Convenient wrapper of
InferConfigMessenger
Given a callable fn that contains Pyro primitive calls and a callable config_fn taking a trace site and returning a dictionary, updates the value of the infer kwarg at a sample site to config_fn(site).
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 config_fn – a callable taking a site and returning an infer dict
Returns: stochastic function decorated with
InferConfigMessenger

lift
(fn=None, *args, **kwargs)¶ Convenient wrapper of
LiftMessenger
Given a stochastic function with param calls and a prior distribution, create a stochastic function where all param calls are replaced by sampling from prior. Prior should be a callable or a dict of names to callables.
Consider the following Pyro program:
>>> def model(x): ... s = pyro.param("s", torch.tensor(0.5)) ... z = pyro.sample("z", dist.Normal(x, s)) ... return z ** 2 >>> lifted_model = pyro.poutine.lift(model, prior={"s": dist.Exponential(0.3)})
lift
makesparam
statements behave likesample
statements using the distributions inprior
. In this example, site s will now behave as if it was replaced withs = pyro.sample("s", dist.Exponential(0.3))
:>>> tr = pyro.poutine.trace(lifted_model).get_trace(0.0) >>> tr.nodes["s"]["type"] == "sample" True >>> tr2 = pyro.poutine.trace(lifted_model).get_trace(0.0) >>> bool((tr2.nodes["s"]["value"] == tr.nodes["s"]["value"]).all()) False
Parameters:  fn – function whose parameters will be lifted to random values
 prior – prior function in the form of a Distribution or a dict of stochastic fns
Returns: fn
decorated with aLiftMessenger

markov
(fn=None, history=1, keep=False, dim=None, name=None)[source]¶ Markov dependency declaration.
This can be used in a variety of ways:  as a context manager  as a decorator for recursive functions  as an iterator for markov chains
Parameters:  history (int) – The number of previous contexts visible from the
current context. Defaults to 1. If zero, this is similar to
pyro.plate
.  keep (bool) – If true, frames are replayable. This is important
when branching: if
keep=True
, neighboring branches at the same level can depend on each other; ifkeep=False
, neighboring branches are independent (conditioned on their share”  dim (int) – An optional dimension to use for this independence index. Interface stub, behavior not yet implemented.
 name (str) – An optional unique name to help inference algorithms match
pyro.markov()
sites between models and guides. Interface stub, behavior not yet implemented.
 history (int) – The number of previous contexts visible from the
current context. Defaults to 1. If zero, this is similar to

mask
(fn=None, *args, **kwargs)¶ Convenient wrapper of
MaskMessenger
Given a stochastic function with some batched sample statements and masking tensor, mask out some of the sample statements elementwise.
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 mask (torch.BoolTensor) – a
{0,1}
valued masking tensor (1 includes a site, 0 excludes a site)
Returns: stochastic function decorated with a
MaskMessenger

queue
(fn=None, queue=None, max_tries=None, extend_fn=None, escape_fn=None, num_samples=None)[source]¶ Used in sequential enumeration over discrete variables.
Given a stochastic function and a queue, return a return value from a complete trace in the queue.
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 queue – a queue data structure like multiprocessing.Queue to hold partial traces
 max_tries – maximum number of attempts to compute a single complete trace
 extend_fn – function (possibly stochastic) that takes a partial trace and a site, and returns a list of extended traces
 escape_fn – function (possibly stochastic) that takes a partial trace and a site, and returns a boolean value to decide whether to exit
 num_samples – optional number of extended traces for extend_fn to return
Returns: stochastic function decorated with poutine logic

reparam
(fn=None, *args, **kwargs)¶ Convenient wrapper of
ReparamMessenger
Reparametrizes each affected sample site into one or more auxiliary sample sites followed by a deterministic transformation [1].
To specify reparameterizers, pass a
config
dict or callable to the constructor. See thepyro.infer.reparam
module for available reparameterizers. [1] Maria I. Gorinova, Dave Moore, Matthew D. Hoffman (2019)
 “Automatic Reparameterisation of Probabilistic Programs” https://arxiv.org/pdf/1906.03028.pdf
Parameters: config (dict or callable) – Configuration, either a dict mapping site name to Reparameterizer
, or a function mapping site toReparameterizer
or None.

replay
(fn=None, *args, **kwargs)¶ Convenient wrapper of
ReplayMessenger
Given a callable that contains Pyro primitive calls, return a callable that runs the original, reusing the values at sites in trace at those sites in the new trace
Consider the following Pyro program:
>>> def model(x): ... s = pyro.param("s", torch.tensor(0.5)) ... z = pyro.sample("z", dist.Normal(x, s)) ... return z ** 2
replay
makessample
statements behave as if they had sampled the values at the corresponding sites in the trace:>>> old_trace = pyro.poutine.trace(model).get_trace(1.0) >>> replayed_model = pyro.poutine.replay(model, trace=old_trace) >>> bool(replayed_model(0.0) == old_trace.nodes["_RETURN"]["value"]) True
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 trace – a
Trace
data structure to replay against  params – dict of names of param sites and constrained values in fn to replay against
Returns: a stochastic function decorated with a
ReplayMessenger

scale
(fn=None, *args, **kwargs)¶ Convenient wrapper of
ScaleMessenger
Given a stochastic function with some sample statements and a positive scale factor, scale the score of all sample and observe sites in the function.
Consider the following Pyro program:
>>> def model(x): ... s = pyro.param("s", torch.tensor(0.5)) ... z = pyro.sample("z", dist.Normal(x, s), obs=1.0) ... return z ** 2
scale
multiplicatively scales the logprobabilities of sample sites:>>> scaled_model = pyro.poutine.scale(model, scale=0.5) >>> scaled_tr = pyro.poutine.trace(scaled_model).get_trace(0.0) >>> unscaled_tr = pyro.poutine.trace(model).get_trace(0.0) >>> bool((scaled_tr.log_prob_sum() == 0.5 * unscaled_tr.log_prob_sum()).all()) True
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 scale – a positive scaling factor
Returns: stochastic function decorated with a
ScaleMessenger

seed
(fn=None, *args, **kwargs)¶ Convenient wrapper of
SeedMessenger
Handler to set the random number generator to a predefined state by setting its seed. This is the same as calling
pyro.set_rng_seed()
before the call to fn. This handler has no additional effect on primitive statements on the standard Pyro backend, but it might interceptpyro.sample
calls in other backends. e.g. the NumPy backend.Parameters:  fn – a stochastic function (callable containing Pyro primitive calls).
 rng_seed (int) – rng seed.

trace
(fn=None, *args, **kwargs)¶ Convenient wrapper of
TraceMessenger
Return a handler that records the inputs and outputs of primitive calls and their dependencies.
Consider the following Pyro program:
>>> def model(x): ... s = pyro.param("s", torch.tensor(0.5)) ... z = pyro.sample("z", dist.Normal(x, s)) ... return z ** 2
We can record its execution using
trace
and use the resulting data structure to compute the logjoint probability of all of the sample sites in the execution or extract all parameters.>>> trace = pyro.poutine.trace(model).get_trace(0.0) >>> logp = trace.log_prob_sum() >>> params = [trace.nodes[name]["value"].unconstrained() for name in trace.param_nodes]
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 graph_type – string that specifies the kind of graph to construct
 param_only – if true, only records params and not samples
Returns: stochastic function decorated with a
TraceMessenger

uncondition
(fn=None, *args, **kwargs)¶ Convenient wrapper of
UnconditionMessenger
Messenger to force the value of observed nodes to be sampled from their distribution, ignoring observations.

config_enumerate
(guide=None, default='parallel', expand=False, num_samples=None, tmc='diagonal')[source]¶ Configures enumeration for all relevant sites in a guide. This is mainly used in conjunction with
TraceEnum_ELBO
.When configuring for exhaustive enumeration of discrete variables, this configures all sample sites whose distribution satisfies
.has_enumerate_support == True
. When configuring for local parallel Monte Carlo sampling viadefault="parallel", num_samples=n
, this configures all sample sites. This does not overwrite existing annotationsinfer={"enumerate": ...}
.This can be used as either a function:
guide = config_enumerate(guide)
or as a decorator:
@config_enumerate def guide1(*args, **kwargs): ... @config_enumerate(default="sequential", expand=True) def guide2(*args, **kwargs): ...
Parameters:  guide (callable) – a pyro model that will be used as a guide in
SVI
.  default (str) – Which enumerate strategy to use, one of “sequential”, “parallel”, or None. Defaults to “parallel”.
 expand (bool) – Whether to expand enumerated sample values. See
enumerate_support()
for details. This only applies to exhaustive enumeration, wherenum_samples=None
. Ifnum_samples
is notNone
, then this samples will always be expanded.  num_samples (int or None) – if not
None
, use local Monte Carlo sampling rather than exhaustive enumeration. This makes sense for both continuous and discrete distributions.  tmc (string or None) – “mixture” or “diagonal” strategies to use in Tensor Monte Carlo
Returns: an annotated guide
Return type: callable
 guide (callable) – a pyro model that will be used as a guide in
Trace¶

class
Trace
(graph_type='flat')[source]¶ Bases:
object
Graph data structure denoting the relationships amongst different pyro primitives in the execution trace.
An execution trace of a Pyro program is a record of every call to
pyro.sample()
andpyro.param()
in a single execution of that program. Traces are directed graphs whose nodes represent primitive calls or input/output, and whose edges represent conditional dependence relationships between those primitive calls. They are created and populated bypoutine.trace
.Each node (or site) in a trace contains the name, input and output value of the site, as well as additional metadata added by inference algorithms or user annotation. In the case of
pyro.sample
, the trace also includes the stochastic function at the site, and any observed data added by users.Consider the following Pyro program:
>>> def model(x): ... s = pyro.param("s", torch.tensor(0.5)) ... z = pyro.sample("z", dist.Normal(x, s)) ... return z ** 2
We can record its execution using
pyro.poutine.trace
and use the resulting data structure to compute the logjoint probability of all of the sample sites in the execution or extract all parameters.>>> trace = pyro.poutine.trace(model).get_trace(0.0) >>> logp = trace.log_prob_sum() >>> params = [trace.nodes[name]["value"].unconstrained() for name in trace.param_nodes]
We can also inspect or manipulate individual nodes in the trace.
trace.nodes
contains acollections.OrderedDict
of site names and metadata corresponding tox
,s
,z
, and the return value:>>> list(name for name in trace.nodes.keys()) # doctest: +SKIP ["_INPUT", "s", "z", "_RETURN"]
Values of
trace.nodes
are dictionaries of node metadata:>>> trace.nodes["z"] # doctest: +SKIP {'type': 'sample', 'name': 'z', 'is_observed': False, 'fn': Normal(), 'value': tensor(0.6480), 'args': (), 'kwargs': {}, 'infer': {}, 'scale': 1.0, 'cond_indep_stack': (), 'done': True, 'stop': False, 'continuation': None}
'infer'
is a dictionary of user or algorithmspecified metadata.'args'
and'kwargs'
are the arguments passed viapyro.sample
tofn.__call__
orfn.log_prob
.'scale'
is used to scale the logprobability of the site when computing the logjoint.'cond_indep_stack'
contains data structures corresponding topyro.plate
contexts appearing in the execution.'done'
,'stop'
, and'continuation'
are only used by Pyro’s internals.Parameters: graph_type (string) – string specifying the kind of trace graph to construct 
add_node
(site_name, **kwargs)[source]¶ Parameters: site_name (string) – the name of the site to be added Adds a site to the trace.
Raises an error when attempting to add a duplicate node instead of silently overwriting.

compute_log_prob
(site_filter=<function Trace.<lambda>>)[source]¶ Compute the sitewise log probabilities of the trace. Each
log_prob
has shape equal to the correspondingbatch_shape
. Eachlog_prob_sum
is a scalar. Both computations are memoized.

compute_score_parts
()[source]¶ Compute the batched local score parts at each site of the trace. Each
log_prob
has shape equal to the correspondingbatch_shape
. Eachlog_prob_sum
is a scalar. All computations are memoized.

edges
¶

format_shapes
(title='Trace Shapes:', last_site=None)[source]¶ Returns a string showing a table of the shapes of all sites in the trace.

log_prob_sum
(site_filter=<function Trace.<lambda>>)[source]¶ Compute the sitewise log probabilities of the trace. Each
log_prob
has shape equal to the correspondingbatch_shape
. Eachlog_prob_sum
is a scalar. The computation oflog_prob_sum
is memoized.Returns: total log probability. Return type: torch.Tensor

nonreparam_stochastic_nodes
¶ Returns: a list of names of sample sites whose stochastic functions are not reparameterizable primitive distributions

observation_nodes
¶ Returns: a list of names of observe sites

pack_tensors
(plate_to_symbol=None)[source]¶ Computes packed representations of tensors in the trace. This should be called after
compute_log_prob()
orcompute_score_parts()
.

param_nodes
¶ Returns: a list of names of param sites

reparameterized_nodes
¶ Returns: a list of names of sample sites whose stochastic functions are reparameterizable primitive distributions

stochastic_nodes
¶ Returns: a list of names of sample sites

Runtime¶

exception
NonlocalExit
(site, *args, **kwargs)[source]¶ Bases:
Exception
Exception for exiting nonlocally from poutine execution.
Used by poutine.EscapeMessenger to return site information.

am_i_wrapped
()[source]¶ Checks whether the current computation is wrapped in a poutine. :returns: bool

apply_stack
(initial_msg)[source]¶ Execute the effect stack at a single site according to the following scheme:
 For each
Messenger
in the stack from bottom to top, executeMessenger._process_message
with the message; if the message field “stop” is True, stop; otherwise, continue  Apply default behavior (
default_process_message
) to finish remaining site execution  For each
Messenger
in the stack from top to bottom, execute_postprocess_message
to update the message and internal messenger state with the site results  If the message field “continuation” is not
None
, call it with the message
Parameters: initial_msg (dict) – the starting version of the trace site Returns: None
 For each

default_process_message
(msg)[source]¶ Default method for processing messages in inference.
Parameters: msg – a message to be processed Returns: None

effectful
(fn=None, type=None)[source]¶ Parameters:  fn – function or callable that performs an effectful computation
 type (str) – the type label of the operation, e.g. “sample”
Wrapper for calling
apply_stack()
to apply any active effects.
Utilities¶

all_escape
(trace, msg)[source]¶ Parameters:  trace – a partial trace
 msg – the message at a Pyro primitive site
Returns: boolean decision value
Utility function that checks if a site is not already in a trace.
Used by EscapeMessenger to decide whether to do a nonlocal exit at a site. Subroutine for approximately integrating out variables for variance reduction.

discrete_escape
(trace, msg)[source]¶ Parameters:  trace – a partial trace
 msg – the message at a Pyro primitive site
Returns: boolean decision value
Utility function that checks if a sample site is discrete and not already in a trace.
Used by EscapeMessenger to decide whether to do a nonlocal exit at a site. Subroutine for integrating out discrete variables for variance reduction.

enum_extend
(trace, msg, num_samples=None)[source]¶ Parameters:  trace – a partial trace
 msg – the message at a Pyro primitive site
 num_samples – maximum number of extended traces to return.
Returns: a list of traces, copies of input trace with one extra site
Utility function to copy and extend a trace with sites based on the input site whose values are enumerated from the support of the input site’s distribution.
Used for exact inference and integrating out discrete variables.

mc_extend
(trace, msg, num_samples=None)[source]¶ Parameters:  trace – a partial trace
 msg – the message at a Pyro primitive site
 num_samples – maximum number of extended traces to return.
Returns: a list of traces, copies of input trace with one extra site
Utility function to copy and extend a trace with sites based on the input site whose values are sampled from the input site’s function.
Used for Monte Carlo marginalization of individual sample sites.
Messengers¶
Messenger objects contain the implementations of the effects exposed by handlers. Advanced users may modify the implementations of messengers behind existing handlers or write new messengers that implement new effects and compose correctly with the rest of the library.
Messenger¶

class
Messenger
[source]¶ Bases:
object
Context manager class that modifies behavior and adds side effects to stochastic functions i.e. callables containing Pyro primitive statements.
This is the base Messenger class. It implements the default behavior for all Pyro primitives, so that the joint distribution induced by a stochastic function fn is identical to the joint distribution induced by
Messenger()(fn)
.Class of transformers for messages passed during inference. Most inference operations are implemented in subclasses of this.

classmethod
register
(fn=None, type=None, post=None)[source]¶ Parameters:  fn – function implementing operation
 type (str) – name of the operation
(also passed to
effectful()
)  post (bool) – if True, use this operation as postprocess
Dynamically add operations to an effect. Useful for generating wrappers for libraries.
Example:
@SomeMessengerClass.register def some_function(msg) ...do_something... return msg

classmethod
unregister
(fn=None, type=None)[source]¶ Parameters:  fn – function implementing operation
 type (str) – name of the operation
(also passed to
effectful()
)
Dynamically remove operations from an effect. Useful for removing wrappers from libraries.
Example:
SomeMessengerClass.unregister(some_function, "name")

classmethod
BlockMessenger¶

class
BlockMessenger
(hide_fn=None, expose_fn=None, hide_all=True, expose_all=False, hide=None, expose=None, hide_types=None, expose_types=None)[source]¶ Bases:
pyro.poutine.messenger.Messenger
This handler selectively hides Pyro primitive sites from the outside world. Default behavior: block everything.
A site is hidden if at least one of the following holds:
hide_fn(msg) is True
or(not expose_fn(msg)) is True
msg["name"] in hide
msg["type"] in hide_types
msg["name"] not in expose and msg["type"] not in expose_types
hide
,hide_types
, andexpose_types
are allNone
For example, suppose the stochastic function fn has two sample sites “a” and “b”. Then any effect outside of
BlockMessenger(fn, hide=["a"])
will not be applied to site “a” and will only see site “b”:>>> def fn(): ... a = pyro.sample("a", dist.Normal(0., 1.)) ... return pyro.sample("b", dist.Normal(a, 1.)) >>> fn_inner = pyro.poutine.trace(fn) >>> fn_outer = pyro.poutine.trace(pyro.poutine.block(fn_inner, hide=["a"])) >>> trace_inner = fn_inner.get_trace() >>> trace_outer = fn_outer.get_trace() >>> "a" in trace_inner True >>> "a" in trace_outer False >>> "b" in trace_inner True >>> "b" in trace_outer True
Parameters:  fn – a stochastic function (callable containing Pyro primitive calls)
 hide_fn – function that takes a site and returns True to hide the site or False/None to expose it. If specified, all other parameters are ignored. Only specify one of hide_fn or expose_fn, not both.
 expose_fn – function that takes a site and returns True to expose the site or False/None to hide it. If specified, all other parameters are ignored. Only specify one of hide_fn or expose_fn, not both.
 hide_all (bool) – hide all sites
 expose_all (bool) – expose all sites normally
 hide (list) – list of site names to hide
 expose (list) – list of site names to be exposed while all others hidden
 hide_types (list) – list of site types to be hidden
 expose_types (lits) – list of site types to be exposed while all others hidden
Returns: stochastic function decorated with a
BlockMessenger
BroadcastMessenger¶

class
BroadcastMessenger
[source]¶ Bases:
pyro.poutine.messenger.Messenger
Automatically broadcasts the batch shape of the stochastic function at a sample site when inside a single or nested plate context. The existing batch_shape must be broadcastable with the size of the
plate
contexts installed in the cond_indep_stack.Notice how model_automatic_broadcast below automates expanding of distribution batch shapes. This makes it easy to modularize a Pyro model as the subcomponents are agnostic of the wrapping
plate
contexts.>>> def model_broadcast_by_hand(): ... with IndepMessenger("batch", 100, dim=2): ... with IndepMessenger("components", 3, dim=1): ... sample = pyro.sample("sample", dist.Bernoulli(torch.ones(3) * 0.5) ... .expand_by(100)) ... assert sample.shape == torch.Size((100, 3)) ... return sample