# Distributions¶

## Base Distribution¶

### Distribution¶

class Distribution(*args, **kwargs)[source]

Bases: object

Base class for probability distributions in NumPyro. The design largely follows from torch.distributions.

Parameters
• batch_shape – The batch shape for the distribution. This designates independent (possibly non-identical) dimensions of a sample from the distribution. This is fixed for a distribution instance and is inferred from the shape of the distribution parameters.

• event_shape – The event shape for the distribution. This designates the dependent dimensions of a sample from the distribution. These are collapsed when we evaluate the log probability density of a batch of samples using .log_prob.

• validate_args – Whether to enable validation of distribution parameters and arguments to .log_prob method.

As an example:

>>> import jax.numpy as jnp
>>> import numpyro.distributions as dist
>>> d = dist.Dirichlet(jnp.ones((2, 3, 4)))
>>> d.batch_shape
(2, 3)
>>> d.event_shape
(4,)

arg_constraints = {}
support = None
has_enumerate_support = False
reparametrized_params = []
tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]
static set_default_validate_args(value)[source]
property batch_shape

Returns the shape over which the distribution parameters are batched.

Returns

batch shape of the distribution.

Return type

tuple

property event_shape

Returns the shape of a single sample from the distribution without batching.

Returns

event shape of the distribution.

Return type

tuple

property event_dim
Returns

Number of dimensions of individual events.

Return type

int

property has_rsample
rsample(key, sample_shape=())[source]
shape(sample_shape=())[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 (tuple) – the size of the iid batch to be drawn from the distribution.

Returns

shape of samples.

Return type

tuple

sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

sample_with_intermediates(key, sample_shape=())[source]

Same as sample except that any intermediate computations are returned (useful for TransformedDistribution).

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(value)[source]

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

to_event(reinterpreted_batch_ndims=None)[source]

Interpret the rightmost reinterpreted_batch_ndims batch dimensions as dependent event dimensions.

Parameters

reinterpreted_batch_ndims – Number of rightmost batch dims to interpret as event dims.

Returns

An instance of Independent distribution.

Return type

numpyro.distributions.distribution.Independent

enumerate_support(expand=True)[source]

Returns an array with shape len(support) x batch_shape containing all values in the support.

expand(batch_shape)[source]

Returns a new ExpandedDistribution instance with batch dimensions expanded to batch_shape.

Parameters

batch_shape (tuple) – batch shape to expand to.

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 its batch_shape. To expand internal dims of self.batch_shape from 1 to something larger, use expand() instead.

Parameters

sample_shape (tuple) – The size of the iid batch to be drawn from the distribution.

Returns

An expanded version of this distribution.

Return type

ExpandedDistribution

Masks a distribution by a boolean or boolean-valued array that is broadcastable to the distributions Distribution.batch_shape .

Parameters

mask (bool or jnp.ndarray) – A boolean or boolean valued array (True includes a site, False excludes a site).

Returns

A masked copy of this distribution.

Return type

MaskedDistribution

Example:

>>> from jax import random
>>> import jax.numpy as jnp
>>> import numpyro
>>> import numpyro.distributions as dist
>>> from numpyro.distributions import constraints
>>> from numpyro.infer import SVI, Trace_ELBO

>>> def model(data, m):
...     f = numpyro.sample("latent_fairness", dist.Beta(1, 1))
...     with numpyro.plate("N", data.shape[0]):
...         # only take into account the values selected by the mask

>>> def guide(data, m):
...     alpha_q = numpyro.param("alpha_q", 5., constraint=constraints.positive)
...     beta_q = numpyro.param("beta_q", 5., constraint=constraints.positive)
...     numpyro.sample("latent_fairness", dist.Beta(alpha_q, beta_q))

>>> data = jnp.concatenate([jnp.ones(5), jnp.zeros(5)])
>>> # select values equal to one
>>> masked_array = jnp.where(data == 1, True, False)
>>> svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
>>> svi_result = svi.run(random.PRNGKey(0), 300, data, masked_array)
>>> params = svi_result.params
>>> # inferred_mean is closer to 1
>>> inferred_mean = params["alpha_q"] / (params["alpha_q"] + params["beta_q"])

classmethod infer_shapes(*args, **kwargs)[source]

Infers batch_shape and event_shape given shapes of args to __init__().

Note

This assumes distribution shape depends only on the shapes of tensor inputs, not in the data contained in those inputs.

Parameters
• *args – Positional args replacing each input arg with a tuple representing the sizes of each tensor input.

• **kwargs – Keywords mapping name of input arg to tuple representing the sizes of each tensor input.

Returns

A pair (batch_shape, event_shape) of the shapes of a distribution that would be created with input args of the given shapes.

Return type

tuple

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

property is_discrete

### ExpandedDistribution¶

class ExpandedDistribution(*args, **kwargs)[source]
arg_constraints = {}
property has_enumerate_support

bool(x) -> bool

Returns True when the argument x is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

property has_rsample
rsample(key, sample_shape=())[source]
property support
sample_with_intermediates(key, sample_shape=())[source]

Same as sample except that any intermediate computations are returned (useful for TransformedDistribution).

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(value)[source]

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

enumerate_support(expand=True)[source]

Returns an array with shape len(support) x batch_shape containing all values in the support.

property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### FoldedDistribution¶

class FoldedDistribution(*args, **kwargs)[source]

Equivalent to TransformedDistribution(base_dist, AbsTransform()), but additionally supports log_prob() .

Parameters

base_dist (Distribution) – A univariate distribution to reflect.

support = <numpyro.distributions.constraints._GreaterThan object>
log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### ImproperUniform¶

class ImproperUniform(*args, **kwargs)[source]

A helper distribution with zero log_prob() over the support domain.

Note

sample method is not implemented for this distribution. In autoguide and mcmc, initial parameters for improper sites are derived from init_to_uniform or init_to_value strategies.

Usage:

>>> from numpyro import sample
>>> from numpyro.distributions import ImproperUniform, Normal, constraints
>>>
>>> def model():
...     # ordered vector with length 10
...     x = sample('x', ImproperUniform(constraints.ordered_vector, (), event_shape=(10,)))
...
...     # real matrix with shape (3, 4)
...     y = sample('y', ImproperUniform(constraints.real, (), event_shape=(3, 4)))
...
...     # a shape-(6, 8) batch of length-5 vectors greater than 3
...     z = sample('z', ImproperUniform(constraints.greater_than(3), (6, 8), event_shape=(5,)))


If you want to set improper prior over all values greater than a, where a is another random variable, you might use

>>> def model():
...     a = sample('a', Normal(0, 1))
...     x = sample('x', ImproperUniform(constraints.greater_than(a), (), event_shape=()))


or if you want to reparameterize it

>>> from numpyro.distributions import TransformedDistribution, transforms
>>> from numpyro.handlers import reparam
>>> from numpyro.infer.reparam import TransformReparam
>>>
>>> def model():
...     a = sample('a', Normal(0, 1))
...     with reparam(config={'x': TransformReparam()}):
...         x = sample('x',
...                    TransformedDistribution(ImproperUniform(constraints.positive, (), ()),
...                                            transforms.AffineTransform(a, 1)))

Parameters
• support (Constraint) – the support of this distribution.

• batch_shape (tuple) – batch shape of this distribution. It is usually safe to set batch_shape=().

• event_shape (tuple) – event shape of this distribution.

arg_constraints = {}
support = <numpyro.distributions.constraints._Dependent object>
log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

tree_flatten()[source]

### Independent¶

class Independent(*args, **kwargs)[source]

Reinterprets batch dimensions of a distribution as event dims by shifting the batch-event dim boundary further to the left.

From a practical standpoint, this is useful when changing the result of log_prob(). For example, a univariate Normal distribution can be interpreted as a multivariate Normal with diagonal covariance:

>>> import numpyro.distributions as dist
>>> normal = dist.Normal(jnp.zeros(3), jnp.ones(3))
>>> [normal.batch_shape, normal.event_shape]
[(3,), ()]
>>> diag_normal = dist.Independent(normal, 1)
>>> [diag_normal.batch_shape, diag_normal.event_shape]
[(), (3,)]

Parameters
• base_distribution (numpyro.distribution.Distribution) – a distribution instance.

• reinterpreted_batch_ndims (int) – the number of batch dims to reinterpret as event dims.

arg_constraints = {}
property support
property has_enumerate_support

bool(x) -> bool

Returns True when the argument x is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

property reparameterized_params
property mean

Mean of the distribution.

property variance

Variance of the distribution.

property has_rsample
rsample(key, sample_shape=())[source]
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(value)[source]

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

expand(batch_shape)[source]

Returns a new ExpandedDistribution instance with batch dimensions expanded to batch_shape.

Parameters

batch_shape (tuple) – batch shape to expand to.

Returns

an instance of ExpandedDistribution.

Return type

ExpandedDistribution

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

Masks a distribution by a boolean array that is broadcastable to the distribution’s Distribution.batch_shape. In the special case mask is False, computation of log_prob() , is skipped, and constant zero values are returned instead.

Parameters

mask (jnp.ndarray or bool) – A boolean or boolean-valued array.

arg_constraints = {}
property has_enumerate_support

bool(x) -> bool

Returns True when the argument x is true, False otherwise. The builtins True and False are the only two instances of the class bool. The class bool is a subclass of the class int, and cannot be subclassed.

property has_rsample
rsample(key, sample_shape=())[source]
property support
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(value)[source]

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

enumerate_support(expand=True)[source]

Returns an array with shape len(support) x batch_shape containing all values in the support.

property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### TransformedDistribution¶

class TransformedDistribution(*args, **kwargs)[source]

Returns a distribution instance obtained as a result of applying a sequence of transforms to a base distribution. For an example, see LogNormal and HalfNormal.

Parameters
• base_distribution – the base distribution over which to apply transforms.

• transforms – a single transform or a list of transforms.

• validate_args – Whether to enable validation of distribution parameters and arguments to .log_prob method.

arg_constraints = {}
property has_rsample
rsample(key, sample_shape=())[source]
property support
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

sample_with_intermediates(key, sample_shape=())[source]

Same as sample except that any intermediate computations are returned (useful for TransformedDistribution).

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]

### Delta¶

class Delta(*args, **kwargs)[source]
arg_constraints = {'log_density': <numpyro.distributions.constraints._Real object>, 'v': <numpyro.distributions.constraints._Dependent object>}
reparametrized_params = ['v', 'log_density']
property support
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### Unit¶

class Unit(*args, **kwargs)[source]

Trivial nonnormalized distribution representing the unit type.

The unit type has a single value with no data, i.e. value.size == 0.

This is used for numpyro.factor() statements.

arg_constraints = {'log_factor': <numpyro.distributions.constraints._Real object>}
support = <numpyro.distributions.constraints._Real object>
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(value)[source]

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

## Continuous Distributions¶

### Beta¶

class Beta(*args, **kwargs)[source]
arg_constraints = {'concentration0': <numpyro.distributions.constraints._GreaterThan object>, 'concentration1': <numpyro.distributions.constraints._GreaterThan object>}
reparametrized_params = ['concentration1', 'concentration0']
support = <numpyro.distributions.constraints._Interval object>
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

### BetaProportion¶

class BetaProportion(*args, **kwargs)[source]

The BetaProportion distribution is a reparameterization of the conventional Beta distribution in terms of a the variate mean and a precision parameter.

Reference:
Beta regression for modelling rates and proportion, Ferrari Silvia, and

Francisco Cribari-Neto. Journal of Applied Statistics 31.7 (2004): 799-815.

arg_constraints = {'concentration': <numpyro.distributions.constraints._GreaterThan object>, 'mean': <numpyro.distributions.constraints._Interval object>}
reparametrized_params = ['mean', 'concentration']
support = <numpyro.distributions.constraints._Interval object>

### Cauchy¶

class Cauchy(*args, **kwargs)[source]
arg_constraints = {'loc': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._Real object>
reparametrized_params = ['loc', 'scale']
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

### Chi2¶

class Chi2(*args, **kwargs)[source]
arg_constraints = {'df': <numpyro.distributions.constraints._GreaterThan object>}
reparametrized_params = ['df']

### Dirichlet¶

class Dirichlet(*args, **kwargs)[source]
arg_constraints = {'concentration': <numpyro.distributions.constraints._IndependentConstraint object>}
reparametrized_params = ['concentration']
support = <numpyro.distributions.constraints._Simplex object>
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

static infer_shapes(concentration)[source]

Infers batch_shape and event_shape given shapes of args to __init__().

Note

This assumes distribution shape depends only on the shapes of tensor inputs, not in the data contained in those inputs.

Parameters
• *args – Positional args replacing each input arg with a tuple representing the sizes of each tensor input.

• **kwargs – Keywords mapping name of input arg to tuple representing the sizes of each tensor input.

Returns

A pair (batch_shape, event_shape) of the shapes of a distribution that would be created with input args of the given shapes.

Return type

tuple

### Exponential¶

class Exponential(*args, **kwargs)[source]
reparametrized_params = ['rate']
arg_constraints = {'rate': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._GreaterThan object>
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

### Gamma¶

class Gamma(*args, **kwargs)[source]
arg_constraints = {'concentration': <numpyro.distributions.constraints._GreaterThan object>, 'rate': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._GreaterThan object>
reparametrized_params = ['concentration', 'rate']
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

### Gumbel¶

class Gumbel(*args, **kwargs)[source]
arg_constraints = {'loc': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._Real object>
reparametrized_params = ['loc', 'scale']
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

### GaussianRandomWalk¶

class GaussianRandomWalk(*args, **kwargs)[source]
arg_constraints = {'scale': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._IndependentConstraint object>
reparametrized_params = ['scale']
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### HalfCauchy¶

class HalfCauchy(*args, **kwargs)[source]
reparametrized_params = ['scale']
support = <numpyro.distributions.constraints._GreaterThan object>
arg_constraints = {'scale': <numpyro.distributions.constraints._GreaterThan object>}
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

property mean

Mean of the distribution.

property variance

Variance of the distribution.

### HalfNormal¶

class HalfNormal(*args, **kwargs)[source]
reparametrized_params = ['scale']
support = <numpyro.distributions.constraints._GreaterThan object>
arg_constraints = {'scale': <numpyro.distributions.constraints._GreaterThan object>}
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

property mean

Mean of the distribution.

property variance

Variance of the distribution.

### InverseGamma¶

class InverseGamma(*args, **kwargs)[source]

Note

We keep the same notation rate as in Pyro but it plays the role of scale parameter of InverseGamma in literatures (e.g. wikipedia: https://en.wikipedia.org/wiki/Inverse-gamma_distribution)

arg_constraints = {'concentration': <numpyro.distributions.constraints._GreaterThan object>, 'rate': <numpyro.distributions.constraints._GreaterThan object>}
reparametrized_params = ['concentration', 'rate']
support = <numpyro.distributions.constraints._GreaterThan object>
property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]

### Laplace¶

class Laplace(*args, **kwargs)[source]
arg_constraints = {'loc': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._Real object>
reparametrized_params = ['loc', 'scale']
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

### LKJ¶

class LKJ(*args, **kwargs)[source]

LKJ distribution for correlation matrices. The distribution is controlled by concentration parameter $$\eta$$ to make the probability of the correlation matrix $$M$$ propotional to $$\det(M)^{\eta - 1}$$. Because of that, when concentration == 1, we have a uniform distribution over correlation matrices.

When concentration > 1, the distribution favors samples with large large determinent. This is useful when we know a priori that the underlying variables are not correlated.

When concentration < 1, the distribution favors samples with small determinent. This is useful when we know a priori that some underlying variables are correlated.

Sample code for using LKJ in the context of multivariate normal sample:

def model(y):  # y has dimension N x d
d = y.shape[1]
N = y.shape[0]
# Vector of variances for each of the d variables
theta = numpyro.sample("theta", dist.HalfCauchy(jnp.ones(d)))

concentration = jnp.ones(1)  # Implies a uniform distribution over correlation matrices
corr_mat = numpyro.sample("corr_mat", dist.LKJ(d, concentration))
sigma = jnp.sqrt(theta)
# we can also use a faster formula cov_mat = jnp.outer(theta, theta) * corr_mat
cov_mat = jnp.matmul(jnp.matmul(jnp.diag(sigma), corr_mat), jnp.diag(sigma))

# Vector of expectations
mu = jnp.zeros(d)

with numpyro.plate("observations", N):
obs = numpyro.sample("obs", dist.MultivariateNormal(mu, covariance_matrix=cov_mat), obs=y)
return obs

Parameters
• dimension (int) – dimension of the matrices

• concentration (ndarray) – concentration/shape parameter of the distribution (often referred to as eta)

• sample_method (str) – Either “cvine” or “onion”. Both methods are proposed in [1] and offer the same distribution over correlation matrices. But they are different in how to generate samples. Defaults to “onion”.

References

[1] Generating random correlation matrices based on vines and extended onion method, Daniel Lewandowski, Dorota Kurowicka, Harry Joe

arg_constraints = {'concentration': <numpyro.distributions.constraints._GreaterThan object>}
reparametrized_params = ['concentration']
support = <numpyro.distributions.constraints._CorrMatrix object>
property mean

Mean of the distribution.

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### LKJCholesky¶

class LKJCholesky(*args, **kwargs)[source]

LKJ distribution for lower Cholesky factors of correlation matrices. The distribution is controlled by concentration parameter $$\eta$$ to make the probability of the correlation matrix $$M$$ generated from a Cholesky factor propotional to $$\det(M)^{\eta - 1}$$. Because of that, when concentration == 1, we have a uniform distribution over Cholesky factors of correlation matrices.

When concentration > 1, the distribution favors samples with large diagonal entries (hence large determinent). This is useful when we know a priori that the underlying variables are not correlated.

When concentration < 1, the distribution favors samples with small diagonal entries (hence small determinent). This is useful when we know a priori that some underlying variables are correlated.

Sample code for using LKJCholesky in the context of multivariate normal sample:

def model(y):  # y has dimension N x d
d = y.shape[1]
N = y.shape[0]
# Vector of variances for each of the d variables
theta = numpyro.sample("theta", dist.HalfCauchy(jnp.ones(d)))
# Lower cholesky factor of a correlation matrix
concentration = jnp.ones(1)  # Implies a uniform distribution over correlation matrices
L_omega = numpyro.sample("L_omega", dist.LKJCholesky(d, concentration))
# Lower cholesky factor of the covariance matrix
sigma = jnp.sqrt(theta)
# we can also use a faster formula L_Omega = sigma[..., None] * L_omega
L_Omega = jnp.matmul(jnp.diag(sigma), L_omega)

# Vector of expectations
mu = jnp.zeros(d)

with numpyro.plate("observations", N):
obs = numpyro.sample("obs", dist.MultivariateNormal(mu, scale_tril=L_Omega), obs=y)
return obs

Parameters
• dimension (int) – dimension of the matrices

• concentration (ndarray) – concentration/shape parameter of the distribution (often referred to as eta)

• sample_method (str) – Either “cvine” or “onion”. Both methods are proposed in [1] and offer the same distribution over correlation matrices. But they are different in how to generate samples. Defaults to “onion”.

References

[1] Generating random correlation matrices based on vines and extended onion method, Daniel Lewandowski, Dorota Kurowicka, Harry Joe

arg_constraints = {'concentration': <numpyro.distributions.constraints._GreaterThan object>}
reparametrized_params = ['concentration']
support = <numpyro.distributions.constraints._CorrCholesky object>
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### LogNormal¶

class LogNormal(*args, **kwargs)[source]
arg_constraints = {'loc': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._GreaterThan object>
reparametrized_params = ['loc', 'scale']
property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]

### Logistic¶

class Logistic(*args, **kwargs)[source]
arg_constraints = {'loc': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._Real object>
reparametrized_params = ['loc', 'scale']
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

### MultivariateNormal¶

class MultivariateNormal(*args, **kwargs)[source]
arg_constraints = {'covariance_matrix': <numpyro.distributions.constraints._PositiveDefinite object>, 'loc': <numpyro.distributions.constraints._IndependentConstraint object>, 'precision_matrix': <numpyro.distributions.constraints._PositiveDefinite object>, 'scale_tril': <numpyro.distributions.constraints._LowerCholesky object>}
support = <numpyro.distributions.constraints._IndependentConstraint object>
reparametrized_params = ['loc', 'covariance_matrix', 'precision_matrix', 'scale_tril']
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

covariance_matrix()[source]
precision_matrix()[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]
static infer_shapes(loc=(), covariance_matrix=None, precision_matrix=None, scale_tril=None)[source]

Infers batch_shape and event_shape given shapes of args to __init__().

Note

This assumes distribution shape depends only on the shapes of tensor inputs, not in the data contained in those inputs.

Parameters
• *args – Positional args replacing each input arg with a tuple representing the sizes of each tensor input.

• **kwargs – Keywords mapping name of input arg to tuple representing the sizes of each tensor input.

Returns

A pair (batch_shape, event_shape) of the shapes of a distribution that would be created with input args of the given shapes.

Return type

tuple

### LowRankMultivariateNormal¶

class LowRankMultivariateNormal(*args, **kwargs)[source]
arg_constraints = {'cov_diag': <numpyro.distributions.constraints._IndependentConstraint object>, 'cov_factor': <numpyro.distributions.constraints._IndependentConstraint object>, 'loc': <numpyro.distributions.constraints._IndependentConstraint object>}
support = <numpyro.distributions.constraints._IndependentConstraint object>
reparametrized_params = ['loc', 'cov_factor', 'cov_diag']
property mean

Mean of the distribution.

variance()[source]

Variance of the distribution.

scale_tril()[source]
covariance_matrix()[source]
precision_matrix()[source]
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

entropy()[source]
static infer_shapes(loc, cov_factor, cov_diag)[source]

Infers batch_shape and event_shape given shapes of args to __init__().

Note

This assumes distribution shape depends only on the shapes of tensor inputs, not in the data contained in those inputs.

Parameters
• *args – Positional args replacing each input arg with a tuple representing the sizes of each tensor input.

• **kwargs – Keywords mapping name of input arg to tuple representing the sizes of each tensor input.

Returns

A pair (batch_shape, event_shape) of the shapes of a distribution that would be created with input args of the given shapes.

Return type

tuple

### Normal¶

class Normal(*args, **kwargs)[source]
arg_constraints = {'loc': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._Real object>
reparametrized_params = ['loc', 'scale']
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

property mean

Mean of the distribution.

property variance

Variance of the distribution.

### Pareto¶

class Pareto(*args, **kwargs)[source]
arg_constraints = {'alpha': <numpyro.distributions.constraints._GreaterThan object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
reparametrized_params = ['scale', 'alpha']
property mean

Mean of the distribution.

property variance

Variance of the distribution.

property support
cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

tree_flatten()[source]

### SoftLaplace¶

class SoftLaplace(*args, **kwargs)[source]

Smooth distribution with Laplace-like tail behavior.

This distribution corresponds to the log-convex density:

z = (value - loc) / scale
log_prob = log(2 / pi) - log(scale) - logaddexp(z, -z)


Like the Laplace density, this density has the heaviest possible tails (asymptotically) while still being log-convex. Unlike the Laplace distribution, this distribution is infinitely differentiable everywhere, and is thus suitable for HMC and Laplace approximation.

Parameters
• loc – Location parameter.

• scale – Scale parameter.

arg_constraints = {'loc': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._Real object>
reparametrized_params = ['loc', 'scale']
log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(value)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

property mean

Mean of the distribution.

property variance

Variance of the distribution.

### StudentT¶

class StudentT(*args, **kwargs)[source]
arg_constraints = {'df': <numpyro.distributions.constraints._GreaterThan object>, 'loc': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._Real object>
reparametrized_params = ['df', 'loc', 'scale']
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(q)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

### Uniform¶

class Uniform(*args, **kwargs)[source]
arg_constraints = {'high': <numpyro.distributions.constraints._Dependent object>, 'low': <numpyro.distributions.constraints._Dependent object>}
reparametrized_params = ['low', 'high']
property support
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

icdf(value)[source]

The inverse cumulative distribution function of this distribution.

Parameters

q – quantile values, should belong to [0, 1].

Returns

the samples whose cdf values equals to q.

property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]
static infer_shapes(low=(), high=())[source]

Infers batch_shape and event_shape given shapes of args to __init__().

Note

This assumes distribution shape depends only on the shapes of tensor inputs, not in the data contained in those inputs.

Parameters
• *args – Positional args replacing each input arg with a tuple representing the sizes of each tensor input.

• **kwargs – Keywords mapping name of input arg to tuple representing the sizes of each tensor input.

Returns

A pair (batch_shape, event_shape) of the shapes of a distribution that would be created with input args of the given shapes.

Return type

tuple

### Weibull¶

class Weibull(*args, **kwargs)[source]
arg_constraints = {'concentration': <numpyro.distributions.constraints._GreaterThan object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._GreaterThan object>
reparametrized_params = ['scale', 'concentration']
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

property mean

Mean of the distribution.

property variance

Variance of the distribution.

## Discrete Distributions¶

### Bernoulli¶

Bernoulli(probs=None, logits=None, validate_args=None)[source]

### BernoulliLogits¶

class BernoulliLogits(*args, **kwargs)[source]
arg_constraints = {'logits': <numpyro.distributions.constraints._Real object>}
support = <numpyro.distributions.constraints._Boolean object>
has_enumerate_support = True
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

probs()[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

enumerate_support(expand=True)[source]

Returns an array with shape len(support) x batch_shape containing all values in the support.

### BernoulliProbs¶

class BernoulliProbs(*args, **kwargs)[source]
arg_constraints = {'probs': <numpyro.distributions.constraints._Interval object>}
support = <numpyro.distributions.constraints._Boolean object>
has_enumerate_support = True
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

logits()[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

enumerate_support(expand=True)[source]

Returns an array with shape len(support) x batch_shape containing all values in the support.

### BetaBinomial¶

class BetaBinomial(*args, **kwargs)[source]

Compound distribution comprising of a beta-binomial pair. The probability of success (probs for the Binomial distribution) is unknown and randomly drawn from a Beta distribution prior to a certain number of Bernoulli trials given by total_count.

Parameters
• concentration1 (numpy.ndarray) – 1st concentration parameter (alpha) for the Beta distribution.

• concentration0 (numpy.ndarray) – 2nd concentration parameter (beta) for the Beta distribution.

• total_count (numpy.ndarray) – number of Bernoulli trials.

arg_constraints = {'concentration0': <numpyro.distributions.constraints._GreaterThan object>, 'concentration1': <numpyro.distributions.constraints._GreaterThan object>, 'total_count': <numpyro.distributions.constraints._IntegerGreaterThan object>}
has_enumerate_support = True
enumerate_support(expand=True)

Returns an array with shape len(support) x batch_shape containing all values in the support.

sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

property support

### Binomial¶

Binomial(total_count=1, probs=None, logits=None, validate_args=None)[source]

### BinomialLogits¶

class BinomialLogits(*args, **kwargs)[source]
arg_constraints = {'logits': <numpyro.distributions.constraints._Real object>, 'total_count': <numpyro.distributions.constraints._IntegerGreaterThan object>}
has_enumerate_support = True
enumerate_support(expand=True)

Returns an array with shape len(support) x batch_shape containing all values in the support.

sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

probs()[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

property support

### BinomialProbs¶

class BinomialProbs(*args, **kwargs)[source]
arg_constraints = {'probs': <numpyro.distributions.constraints._Interval object>, 'total_count': <numpyro.distributions.constraints._IntegerGreaterThan object>}
has_enumerate_support = True
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

logits()[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

property support
enumerate_support(expand=True)[source]

Returns an array with shape len(support) x batch_shape containing all values in the support.

### Categorical¶

Categorical(probs=None, logits=None, validate_args=None)[source]

### CategoricalLogits¶

class CategoricalLogits(*args, **kwargs)[source]
arg_constraints = {'logits': <numpyro.distributions.constraints._IndependentConstraint object>}
has_enumerate_support = True
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

probs()[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

property support
enumerate_support(expand=True)[source]

Returns an array with shape len(support) x batch_shape containing all values in the support.

### CategoricalProbs¶

class CategoricalProbs(*args, **kwargs)[source]
arg_constraints = {'probs': <numpyro.distributions.constraints._Simplex object>}
has_enumerate_support = True
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

logits()[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

property support
enumerate_support(expand=True)[source]

Returns an array with shape len(support) x batch_shape containing all values in the support.

### DirichletMultinomial¶

class DirichletMultinomial(*args, **kwargs)[source]

Compound distribution comprising of a dirichlet-multinomial pair. The probability of classes (probs for the Multinomial distribution) is unknown and randomly drawn from a Dirichlet distribution prior to a certain number of Categorical trials given by total_count.

Parameters
• concentration (numpy.ndarray) – concentration parameter (alpha) for the Dirichlet distribution.

• total_count (numpy.ndarray) – number of Categorical trials.

arg_constraints = {'concentration': <numpyro.distributions.constraints._IndependentConstraint object>, 'total_count': <numpyro.distributions.constraints._IntegerGreaterThan object>}
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

property support
static infer_shapes(concentration, total_count=())[source]

Infers batch_shape and event_shape given shapes of args to __init__().

Note

This assumes distribution shape depends only on the shapes of tensor inputs, not in the data contained in those inputs.

Parameters
• *args – Positional args replacing each input arg with a tuple representing the sizes of each tensor input.

• **kwargs – Keywords mapping name of input arg to tuple representing the sizes of each tensor input.

Returns

A pair (batch_shape, event_shape) of the shapes of a distribution that would be created with input args of the given shapes.

Return type

tuple

### GammaPoisson¶

class GammaPoisson(*args, **kwargs)[source]

Compound distribution comprising of a gamma-poisson pair, also referred to as a gamma-poisson mixture. The rate parameter for the Poisson distribution is unknown and randomly drawn from a Gamma distribution.

Parameters
• concentration (numpy.ndarray) – shape parameter (alpha) of the Gamma distribution.

• rate (numpy.ndarray) – rate parameter (beta) for the Gamma distribution.

arg_constraints = {'concentration': <numpyro.distributions.constraints._GreaterThan object>, 'rate': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._IntegerGreaterThan object>
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

### Geometric¶

Geometric(probs=None, logits=None, validate_args=None)[source]

### GeometricLogits¶

class GeometricLogits(*args, **kwargs)[source]
arg_constraints = {'logits': <numpyro.distributions.constraints._Real object>}
support = <numpyro.distributions.constraints._IntegerGreaterThan object>
probs()[source]
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

### GeometricProbs¶

class GeometricProbs(*args, **kwargs)[source]
arg_constraints = {'probs': <numpyro.distributions.constraints._Interval object>}
support = <numpyro.distributions.constraints._IntegerGreaterThan object>
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

logits()[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

### Multinomial¶

Multinomial(total_count=1, probs=None, logits=None, validate_args=None)[source]

### MultinomialLogits¶

class MultinomialLogits(*args, **kwargs)[source]
arg_constraints = {'logits': <numpyro.distributions.constraints._IndependentConstraint object>, 'total_count': <numpyro.distributions.constraints._IntegerGreaterThan object>}
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

probs()[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

property support
static infer_shapes(logits, total_count)[source]

Infers batch_shape and event_shape given shapes of args to __init__().

Note

This assumes distribution shape depends only on the shapes of tensor inputs, not in the data contained in those inputs.

Parameters
• *args – Positional args replacing each input arg with a tuple representing the sizes of each tensor input.

• **kwargs – Keywords mapping name of input arg to tuple representing the sizes of each tensor input.

Returns

A pair (batch_shape, event_shape) of the shapes of a distribution that would be created with input args of the given shapes.

Return type

tuple

### MultinomialProbs¶

class MultinomialProbs(*args, **kwargs)[source]
arg_constraints = {'probs': <numpyro.distributions.constraints._Simplex object>, 'total_count': <numpyro.distributions.constraints._IntegerGreaterThan object>}
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

logits()[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

property support
static infer_shapes(probs, total_count)[source]

Infers batch_shape and event_shape given shapes of args to __init__().

Note

This assumes distribution shape depends only on the shapes of tensor inputs, not in the data contained in those inputs.

Parameters
• *args – Positional args replacing each input arg with a tuple representing the sizes of each tensor input.

• **kwargs – Keywords mapping name of input arg to tuple representing the sizes of each tensor input.

Returns

A pair (batch_shape, event_shape) of the shapes of a distribution that would be created with input args of the given shapes.

Return type

tuple

### OrderedLogistic¶

class OrderedLogistic(*args, **kwargs)[source]

A categorical distribution with ordered outcomes.

References:

1. Stan Functions Reference, v2.20 section 12.6, Stan Development Team

Parameters
• predictor (numpy.ndarray) – prediction in real domain; typically this is output of a linear model.

• cutpoints (numpy.ndarray) – positions in real domain to separate categories.

arg_constraints = {'cutpoints': <numpyro.distributions.constraints._OrderedVector object>, 'predictor': <numpyro.distributions.constraints._Real object>}
static infer_shapes(predictor, cutpoints)[source]

Infers batch_shape and event_shape given shapes of args to __init__().

Note

This assumes distribution shape depends only on the shapes of tensor inputs, not in the data contained in those inputs.

Parameters
• *args – Positional args replacing each input arg with a tuple representing the sizes of each tensor input.

• **kwargs – Keywords mapping name of input arg to tuple representing the sizes of each tensor input.

Returns

A pair (batch_shape, event_shape) of the shapes of a distribution that would be created with input args of the given shapes.

Return type

tuple

### NegativeBinomial¶

NegativeBinomial(total_count, probs=None, logits=None, validate_args=None)[source]

### NegativeBinomialLogits¶

class NegativeBinomialLogits(*args, **kwargs)[source]
arg_constraints = {'logits': <numpyro.distributions.constraints._Real object>, 'total_count': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._IntegerGreaterThan object>
log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

### NegativeBinomialProbs¶

class NegativeBinomialProbs(*args, **kwargs)[source]
arg_constraints = {'probs': <numpyro.distributions.constraints._Interval object>, 'total_count': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._IntegerGreaterThan object>

### NegativeBinomial2¶

class NegativeBinomial2(*args, **kwargs)[source]

Another parameterization of GammaPoisson with rate is replaced by mean.

arg_constraints = {'concentration': <numpyro.distributions.constraints._GreaterThan object>, 'mean': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._IntegerGreaterThan object>

### Poisson¶

class Poisson(*args, **kwargs)[source]

Creates a Poisson distribution parameterized by rate, the rate parameter.

Samples are nonnegative integers, with a pmf given by

$\mathrm{rate}^k \frac{e^{-\mathrm{rate}}}{k!}$
Parameters
arg_constraints = {'rate': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._IntegerGreaterThan object>
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the distribution.

property variance

Variance of the distribution.

cdf(value)[source]

The cummulative distribution function of this distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

### PRNGIdentity¶

class PRNGIdentity(*args, **kwargs)[source]

Distribution over PRNGKey(). This can be used to draw a batch of PRNGKey() using the seed handler. Only sample method is supported.

sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

### ZeroInflatedDistribution¶

ZeroInflatedDistribution(base_dist, *, gate=None, gate_logits=None, validate_args=None)[source]

Generic Zero Inflated distribution.

Parameters
• base_dist (Distribution) – the base distribution.

• gate (numpy.ndarray) – probability of extra zeros given via a Bernoulli distribution.

• gate_logits (numpy.ndarray) – logits of extra zeros given via a Bernoulli distribution.

### ZeroInflatedPoisson¶

class ZeroInflatedPoisson(*args, **kwargs)[source]

Bases: numpyro.distributions.discrete.ZeroInflatedProbs

A Zero Inflated Poisson distribution.

Parameters
arg_constraints = {'gate': <numpyro.distributions.constraints._Interval object>, 'rate': <numpyro.distributions.constraints._GreaterThan object>}
support = <numpyro.distributions.constraints._IntegerGreaterThan object>

### ZeroInflatedNegativeBinomial2¶

ZeroInflatedNegativeBinomial2(mean, concentration, *, gate=None, gate_logits=None, validate_args=None)[source]

## Mixture Distributions¶

### MixtureSameFamily¶

class MixtureSameFamily(*args, **kwargs)[source]

Marginalized Finite Mixture distribution of vectorized components.

The components being a vectorized distribution implies that all components are from the same family, represented by a single Distribution object.

Parameters
• mixing_distribution (numpyro.distribution.Distribution) – The mixing distribution to select the components. Needs to be a categorical.

• component_distribution (numpyro.distribution.Distribution) – Vectorized component distribution.

As an example:

>>> import jax
>>> import jax.numpy as jnp
>>> import numpyro.distributions as dist
>>> mixing_dist = dist.Categorical(probs=jnp.ones(3) / 3.)
>>> component_dist = dist.Normal(loc=jnp.zeros(3), scale=jnp.ones(3))
>>> mixture = dist.MixtureSameFamily(mixing_dist, component_dist)
>>> mixture.sample(jax.random.PRNGKey(42)).shape
()

property mixture_size

Returns the number of distributions in the mixture

Returns

number of mixtures.

Return type

int

property mixing_distribution

Returns the mixing distribution

Returns

Categorical distribution

Return type

Categorical

property mixture_dim
property component_distribution

Return the vectorized distribution of components being mixed.

Returns

Component distribution

Return type

Distribution

property support
property is_discrete
tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]
property mean

Mean of the distribution.

property variance

Variance of the distribution.

cdf(samples)[source]

The cumulative distribution function of this mixture distribution.

Parameters

value – samples from this distribution.

Returns

output of the cummulative distribution function evaluated at value.

Raises

NotImplementedError if the component distribution does not implement the cdf method.

sample_with_intermediates(key, sample_shape=())[source]

Same as sample except that the sampled mixture components are also returned.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

Tuple (samples, indices)

Return type

tuple

sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

## Directional Distributions¶

### ProjectedNormal¶

class ProjectedNormal(*args, **kwargs)[source]

Projected isotropic normal distribution of arbitrary dimension.

This distribution over directional data is qualitatively similar to the von Mises and von Mises-Fisher distributions, but permits tractable variational inference via reparametrized gradients.

To use this distribution with autoguides and HMC, use handlers.reparam with a ProjectedNormalReparam reparametrizer in the model, e.g.:

@handlers.reparam(config={"direction": ProjectedNormalReparam()})
def model():
direction = numpyro.sample("direction",
ProjectedNormal(zeros(3)))
...


Note

This implements log_prob() only for dimensions {2,3}.

[1] D. Hernandez-Stumpfhauser, F.J. Breidt, M.J. van der Woerd (2017)

“The General Projected Normal Distribution of Arbitrary Dimension: Modeling and Bayesian Inference” https://projecteuclid.org/euclid.ba/1453211962

arg_constraints = {'concentration': <numpyro.distributions.constraints._IndependentConstraint object>}
reparametrized_params = ['concentration']
support = <numpyro.distributions.constraints._Sphere object>
property mean

Note this is the mean in the sense of a centroid in the submanifold that minimizes expected squared geodesic distance.

property mode
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(value)[source]

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

static infer_shapes(concentration)[source]

Infers batch_shape and event_shape given shapes of args to __init__().

Note

This assumes distribution shape depends only on the shapes of tensor inputs, not in the data contained in those inputs.

Parameters
• *args – Positional args replacing each input arg with a tuple representing the sizes of each tensor input.

• **kwargs – Keywords mapping name of input arg to tuple representing the sizes of each tensor input.

Returns

A pair (batch_shape, event_shape) of the shapes of a distribution that would be created with input args of the given shapes.

Return type

tuple

### SineBivariateVonMises¶

class SineBivariateVonMises(*args, **kwargs)[source]

Unimodal distribution of two dependent angles on the 2-torus (S^1 ⨂ S^1) given by

$C^{-1}\exp(\kappa_1\cos(x_1-\mu_1) + \kappa_2\cos(x_2 -\mu_2) + \rho\sin(x_1 - \mu_1)\sin(x_2 - \mu_2))$

and

$C = (2\pi)^2 \sum_{i=0} {2i \choose i} \left(\frac{\rho^2}{4\kappa_1\kappa_2}\right)^i I_i(\kappa_1)I_i(\kappa_2),$

where $$I_i(\cdot)$$ is the modified bessel function of first kind, mu’s are the locations of the distribution, kappa’s are the concentration and rho gives the correlation between angles $$x_1$$ and $$x_2$$. This distribution is helpful for modeling coupled angles such as torsion angles in peptide chains.

To infer parameters, use NUTS or HMC with priors that avoid parameterizations where the distribution becomes bimodal; see note below.

Note

Sample efficiency drops as

$\frac{\rho}{\kappa_1\kappa_2} \rightarrow 1$

because the distribution becomes increasingly bimodal.

Note

The correlation and weighted_correlation params are mutually exclusive.

Note

In the context of SVI, this distribution can be used as a likelihood but not for latent variables.

** References: **
1. Probabilistic model for two dependent circular variables Singh, H., Hnizdo, V., and Demchuck, E. (2002)

Parameters
• phi_loc (np.ndarray) – location of first angle

• psi_loc (np.ndarray) – location of second angle

• phi_concentration (np.ndarray) – concentration of first angle

• psi_concentration (np.ndarray) – concentration of second angle

• correlation (np.ndarray) – correlation between the two angles

• weighted_correlation (np.ndarray) – set correlation to weigthed_corr * sqrt(phi_conc*psi_conc) to avoid bimodality (see note).

arg_constraints = {'correlation': <numpyro.distributions.constraints._Real object>, 'phi_concentration': <numpyro.distributions.constraints._GreaterThan object>, 'phi_loc': <numpyro.distributions.constraints._Interval object>, 'psi_concentration': <numpyro.distributions.constraints._GreaterThan object>, 'psi_loc': <numpyro.distributions.constraints._Interval object>}
support = <numpyro.distributions.constraints._IndependentConstraint object>
max_sample_iter = 1000
norm_const()[source]
log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

sample(key, sample_shape=())[source]
** References: **
1. A New Unified Approach for the Simulation of aWide Class of Directional Distributions John T. Kent, Asaad M. Ganeiber & Kanti V. Mardia (2018)

property mean

Computes circular mean of distribution. Note: same as location when mapped to support [-pi, pi]

### SineSkewed¶

class SineSkewed(*args, **kwargs)[source]

Sine-skewing [1] is a procedure for producing a distribution that breaks pointwise symmetry on a torus distribution. The new distribution is called the Sine Skewed X distribution, where X is the name of the (symmetric) base distribution. Torus distributions are distributions with support on products of circles (i.e., ⨂^d S^1 where S^1=[-pi,pi) ). So, a 0-torus is a point, the 1-torus is a circle, and the 2-torus is commonly associated with the donut shape.

The sine skewed X distribution is parameterized by a weight parameter for each dimension of the event of X. For example with a von Mises distribution over a circle (1-torus), the sine skewed von Mises distribution has one skew parameter. The skewness parameters can be inferred using HMC or NUTS. For example, the following will produce a prior over skewness for the 2-torus,:

@numpyro.handlers.reparam(config={'phi_loc': CircularReparam(), 'psi_loc': CircularReparam()})
def model(obs):
# Sine priors
phi_loc = numpyro.sample('phi_loc', VonMises(pi, 2.))
psi_loc = numpyro.sample('psi_loc', VonMises(-pi / 2, 2.))
phi_conc = numpyro.sample('phi_conc', Beta(1., 1.))
psi_conc = numpyro.sample('psi_conc', Beta(1., 1.))
corr_scale = numpyro.sample('corr_scale', Beta(2., 5.))

# Skewing prior
ball_trans = L1BallTransform()
skewness = numpyro.sample('skew_phi', Normal(0, 0.5).expand((2,)))
skewness = ball_trans(skewness)  # constraint sum |skewness_i| <= 1

with numpyro.plate('obs_plate'):
sine = SineBivariateVonMises(phi_loc=phi_loc, psi_loc=psi_loc,
phi_concentration=70 * phi_conc,
psi_concentration=70 * psi_conc,
weighted_correlation=corr_scale)
return numpyro.sample('phi_psi', SineSkewed(sine, skewness), obs=obs)


To ensure the skewing does not alter the normalization constant of the (sine bivariate von Mises) base distribution the skewness parameters are constraint. The constraint requires the sum of the absolute values of skewness to be less than or equal to one. We can use the L1BallTransform to achieve this.

In the context of SVI, this distribution can freely be used as a likelihood, but use as latent variables it will lead to slow inference for 2 and higher dim toruses. This is because the base_dist cannot be reparameterized.

Note

An event in the base distribution must be on a d-torus, so the event_shape must be (d,).

Note

For the skewness parameter, it must hold that the sum of the absolute value of its weights for an event must be less than or equal to one. See eq. 2.1 in [1].

** References: **
1. Sine-skewed toroidal distributions and their application in protein bioinformatics

Ameijeiras-Alonso, J., Ley, C. (2019)

Parameters
• base_dist (numpyro.distributions.Distribution) – base density on a d-dimensional torus. Supported base distributions include: 1D VonMises, SineBivariateVonMises, 1D ProjectedNormal, and Uniform (-pi, pi).

• skewness (jax.numpy.array) – skewness of the distribution.

arg_constraints = {'skewness': <numpyro.distributions.constraints._L1Ball object>}
support = <numpyro.distributions.constraints._IndependentConstraint object>
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(value)[source]

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Mean of the base distribution

### VonMises¶

class VonMises(*args, **kwargs)[source]

The von Mises distribution, also known as the circular normal distribution.

This distribution is supported by a circular constraint from -pi to +pi. By default, the circular support behaves like constraints.interval(-math.pi, math.pi). To avoid issues at the boundaries of this interval during sampling, you should reparameterize this distribution using handlers.reparam with a CircularReparam reparametrizer in the model, e.g.:

@handlers.reparam(config={"direction": CircularReparam()})
def model():
direction = numpyro.sample("direction", VonMises(0.0, 4.0))
...

arg_constraints = {'concentration': <numpyro.distributions.constraints._GreaterThan object>, 'loc': <numpyro.distributions.constraints._Real object>}
reparametrized_params = ['loc']
support = <numpyro.distributions.constraints._Interval object>
sample(key, sample_shape=())[source]

Generate sample from von Mises distribution

Parameters
• key – random number generator key

• sample_shape – shape of samples

Returns

samples from von Mises

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

property mean

Computes circular mean of distribution. NOTE: same as location when mapped to support [-pi, pi]

property variance

Computes circular variance of distribution

## Truncated Distributions¶

### LeftTruncatedDistribution¶

class LeftTruncatedDistribution(*args, **kwargs)[source]
arg_constraints = {'low': <numpyro.distributions.constraints._Real object>}
reparametrized_params = ['low']
supported_types = (<class 'numpyro.distributions.continuous.Cauchy'>, <class 'numpyro.distributions.continuous.Laplace'>, <class 'numpyro.distributions.continuous.Logistic'>, <class 'numpyro.distributions.continuous.Normal'>, <class 'numpyro.distributions.continuous.SoftLaplace'>, <class 'numpyro.distributions.continuous.StudentT'>)
property support
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### RightTruncatedDistribution¶

class RightTruncatedDistribution(*args, **kwargs)[source]
arg_constraints = {'high': <numpyro.distributions.constraints._Real object>}
reparametrized_params = ['high']
supported_types = (<class 'numpyro.distributions.continuous.Cauchy'>, <class 'numpyro.distributions.continuous.Laplace'>, <class 'numpyro.distributions.continuous.Logistic'>, <class 'numpyro.distributions.continuous.Normal'>, <class 'numpyro.distributions.continuous.SoftLaplace'>, <class 'numpyro.distributions.continuous.StudentT'>)
property support
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### TruncatedCauchy¶

class TruncatedCauchy(*args, **kwargs)[source]
arg_constraints = {'loc': <numpyro.distributions.constraints._Real object>, 'low': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
reparametrized_params = ['low', 'loc', 'scale']
property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### TruncatedDistribution¶

TruncatedDistribution(base_dist, low=None, high=None, validate_args=None)[source]

A function to generate a truncated distribution.

Parameters
• base_dist – The base distribution to be truncated. This should be a univariate distribution. Currently, only the following distributions are supported: Cauchy, Laplace, Logistic, Normal, and StudentT.

• low – the value which is used to truncate the base distribution from below. Setting this parameter to None to not truncate from below.

• high – the value which is used to truncate the base distribution from above. Setting this parameter to None to not truncate from above.

### TruncatedNormal¶

class TruncatedNormal(*args, **kwargs)[source]
arg_constraints = {'loc': <numpyro.distributions.constraints._Real object>, 'low': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}
reparametrized_params = ['low', 'loc', 'scale']
property mean

Mean of the distribution.

property variance

Variance of the distribution.

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### TruncatedPolyaGamma¶

class TruncatedPolyaGamma(*args, **kwargs)[source]
truncation_point = 2.5
num_log_prob_terms = 7
num_gamma_variates = 8
arg_constraints = {}
support = <numpyro.distributions.constraints._Interval object>
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

### TwoSidedTruncatedDistribution¶

class TwoSidedTruncatedDistribution(*args, **kwargs)[source]
arg_constraints = {'high': <numpyro.distributions.constraints._Dependent object>, 'low': <numpyro.distributions.constraints._Dependent object>}
reparametrized_params = ['low', 'high']
supported_types = (<class 'numpyro.distributions.continuous.Cauchy'>, <class 'numpyro.distributions.continuous.Laplace'>, <class 'numpyro.distributions.continuous.Logistic'>, <class 'numpyro.distributions.continuous.Normal'>, <class 'numpyro.distributions.continuous.SoftLaplace'>, <class 'numpyro.distributions.continuous.StudentT'>)
property support
sample(key, sample_shape=())[source]

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters
• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns

an array of shape sample_shape + batch_shape + event_shape

Return type

numpy.ndarray

log_prob(*args, **kwargs)

Evaluates the log probability density for a batch of samples given by value.

Parameters

value – A batch of samples from the distribution.

Returns

an array with shape value.shape[:-self.event_shape]

Return type

numpy.ndarray

tree_flatten()[source]
classmethod tree_unflatten(aux_data, params)[source]

## TensorFlow Distributions¶

Thin wrappers around TensorFlow Probability (TFP) distributions. For details on the TFP distribution interface, see its Distribution docs.

## Constraints¶

### Constraint¶

class Constraint[source]

Bases: object

Abstract base class for constraints.

A constraint object represents a region over which a variable is valid, e.g. within which a variable can be optimized.

is_discrete = False
event_dim = 0
check(value)[source]

Returns a byte tensor of sample_shape + batch_shape indicating whether each event in value satisfies this constraint.

feasible_like(prototype)[source]

Get a feasible value which has the same shape as dtype as prototype.

### boolean¶

boolean = <numpyro.distributions.constraints._Boolean object>

### circular¶

circular = <numpyro.distributions.constraints._Interval object>

### corr_cholesky¶

corr_cholesky = <numpyro.distributions.constraints._CorrCholesky object>

### corr_matrix¶

corr_matrix = <numpyro.distributions.constraints._CorrMatrix object>

### dependent¶

dependent = <numpyro.distributions.constraints._Dependent object>

Placeholder for variables whose support depends on other variables. These variables obey no simple coordinate-wise constraints.

Parameters
• is_discrete (bool) – Optional value of .is_discrete in case this can be computed statically. If not provided, access to the .is_discrete attribute will raise a NotImplementedError.

• event_dim (int) – Optional value of .event_dim in case this can be computed statically. If not provided, access to the .event_dim attribute will raise a NotImplementedError.

### greater_than¶

greater_than(lower_bound)

Abstract base class for constraints.

A constraint object represents a region over which a variable is valid, e.g. within which a variable can be optimized.

### integer_interval¶

integer_interval(lower_bound, upper_bound)

Abstract base class for constraints.

A constraint object represents a region over which a variable is valid, e.g. within which a variable can be optimized.

### integer_greater_than¶

integer_greater_than(lower_bound)

Abstract base class for constraints.

A constraint object represents a region over which a variable is valid, e.g. within which a variable can be optimized.

### interval¶

interval(lower_bound, upper_bound)

Abstract base class for constraints.

A constraint object represents a region over which a variable is valid, e.g. within which a variable can be optimized.

### l1_ball¶

l1_ball(x)

Constrain to the L1 ball of any dimension.

### less_than¶

less_than(upper_bound)

Abstract base class for constraints.

A constraint object represents a region over which a variable is valid, e.g. within which a variable can be optimized.

### lower_cholesky¶

lower_cholesky = <numpyro.distributions.constraints._LowerCholesky object>

### multinomial¶

multinomial(upper_bound)

Abstract base class for constraints.

A constraint object represents a region over which a variable is valid, e.g. within which a variable can be optimized.

### nonnegative_integer¶

nonnegative_integer = <numpyro.distributions.constraints._IntegerGreaterThan object>

### ordered_vector¶

ordered_vector = <numpyro.distributions.constraints._OrderedVector object>

### positive¶

positive = <numpyro.distributions.constraints._GreaterThan object>

### positive_definite¶

positive_definite = <numpyro.distributions.constraints._PositiveDefinite object>

### positive_integer¶

positive_integer = <numpyro.distributions.constraints._IntegerGreaterThan object>

### positive_ordered_vector¶

positive_ordered_vector = <numpyro.distributions.constraints._PositiveOrderedVector object>

Constrains to a positive real-valued tensor where the elements are monotonically increasing along the event_shape dimension.

### real¶

real = <numpyro.distributions.constraints._Real object>

### real_vector¶

real_vector = <numpyro.distributions.constraints._IndependentConstraint object>

Wraps a constraint by aggregating over reinterpreted_batch_ndims-many dims in check(), so that an event is valid only if all its independent entries are valid.

### scaled_unit_lower_cholesky¶

scaled_unit_lower_cholesky = <numpyro.distributions.constraints._ScaledUnitLowerCholesky object>

### softplus_positive¶

softplus_positive = <numpyro.distributions.constraints._SoftplusPositive object>

### softplus_lower_cholesky¶

softplus_lower_cholesky = <numpyro.distributions.constraints._SoftplusLowerCholesky object>

### simplex¶

simplex = <numpyro.distributions.constraints._Simplex object>

### sphere¶

sphere = <numpyro.distributions.constraints._Sphere object>

Constrain to the Euclidean sphere of any dimension.

### unit_interval¶

unit_interval = <numpyro.distributions.constraints._Interval object>

## Transforms¶

### biject_to¶

biject_to(constraint)

### Transform¶

class Transform[source]

Bases: object

domain = <numpyro.distributions.constraints._Real object>
codomain = <numpyro.distributions.constraints._Real object>
property event_dim
property inv
log_abs_det_jacobian(x, y, intermediates=None)[source]
call_with_intermediates(x)[source]
forward_shape(shape)[source]

Infers the shape of the forward computation, given the input shape. Defaults to preserving shape.

inverse_shape(shape)[source]

Infers the shapes of the inverse computation, given the output shape. Defaults to preserving shape.

### AbsTransform¶

class AbsTransform[source]
domain = <numpyro.distributions.constraints._Real object>
codomain = <numpyro.distributions.constraints._GreaterThan object>

### AffineTransform¶

class AffineTransform(loc, scale, domain=<numpyro.distributions.constraints._Real object>)[source]

Note

When scale is a JAX tracer, we always assume that scale > 0 when calculating codomain.

property codomain
log_abs_det_jacobian(x, y, intermediates=None)[source]
forward_shape(shape)[source]

Infers the shape of the forward computation, given the input shape. Defaults to preserving shape.

inverse_shape(shape)[source]

Infers the shapes of the inverse computation, given the output shape. Defaults to preserving shape.

### CholeskyTransform¶

class CholeskyTransform[source]

Transform via the mapping $$y = cholesky(x)$$, where x is a positive definite matrix.

domain = <numpyro.distributions.constraints._PositiveDefinite object>
codomain = <numpyro.distributions.constraints._LowerCholesky object>
log_abs_det_jacobian(x, y, intermediates=None)[source]

### ComposeTransform¶

class ComposeTransform(parts)[source]
property domain
property codomain
log_abs_det_jacobian(x, y, intermediates=None)[source]
call_with_intermediates(x)[source]
forward_shape(shape)[source]

Infers the shape of the forward computation, given the input shape. Defaults to preserving shape.

inverse_shape(shape)[source]

Infers the shapes of the inverse computation, given the output shape. Defaults to preserving shape.

### CorrCholeskyTransform¶

class CorrCholeskyTransform[source]

Transforms a uncontrained real vector $$x$$ with length $$D*(D-1)/2$$ into the Cholesky factor of a D-dimension correlation matrix. This Cholesky factor is a lower triangular matrix with positive diagonals and unit Euclidean norm for each row. The transform is processed as follows:

1. First we convert $$x$$ into a lower triangular matrix with the following order:

$\begin{split}\begin{bmatrix} 1 & 0 & 0 & 0 \\ x_0 & 1 & 0 & 0 \\ x_1 & x_2 & 1 & 0 \\ x_3 & x_4 & x_5 & 1 \end{bmatrix}\end{split}$

2. For each row $$X_i$$ of the lower triangular part, we apply a signed version of class StickBreakingTransform to transform $$X_i$$ into a unit Euclidean length vector using the following steps:

1. Scales into the interval $$(-1, 1)$$ domain: $$r_i = \tanh(X_i)$$.

2. Transforms into an unsigned domain: $$z_i = r_i^2$$.

3. Applies $$s_i = StickBreakingTransform(z_i)$$.

4. Transforms back into signed domain: $$y_i = (sign(r_i), 1) * \sqrt{s_i}$$.

domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._CorrCholesky object>
log_abs_det_jacobian(x, y, intermediates=None)[source]
forward_shape(shape)[source]

Infers the shape of the forward computation, given the input shape. Defaults to preserving shape.

inverse_shape(shape)[source]

Infers the shapes of the inverse computation, given the output shape. Defaults to preserving shape.

### CorrMatrixCholeskyTransform¶

class CorrMatrixCholeskyTransform[source]

Transform via the mapping $$y = cholesky(x)$$, where x is a correlation matrix.

domain = <numpyro.distributions.constraints._CorrMatrix object>
codomain = <numpyro.distributions.constraints._CorrCholesky object>
log_abs_det_jacobian(x, y, intermediates=None)[source]

### ExpTransform¶

class ExpTransform(domain=<numpyro.distributions.constraints._Real object>)[source]
property codomain
log_abs_det_jacobian(x, y, intermediates=None)[source]

### IdentityTransform¶

class IdentityTransform[source]
log_abs_det_jacobian(x, y, intermediates=None)[source]

### InvCholeskyTransform¶

class InvCholeskyTransform(domain=<numpyro.distributions.constraints._LowerCholesky object>)[source]

Transform via the mapping $$y = x @ x.T$$, where x is a lower triangular matrix with positive diagonal.

property codomain
log_abs_det_jacobian(x, y, intermediates=None)[source]

### L1BallTransform¶

class L1BallTransform[source]

Transforms a uncontrained real vector $$x$$ into the unit L1 ball.

domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._L1Ball object>
log_abs_det_jacobian(x, y, intermediates=None)[source]

### LowerCholeskyAffine¶

class LowerCholeskyAffine(loc, scale_tril)[source]

Transform via the mapping $$y = loc + scale\_tril\ @\ x$$.

Parameters
• loc – a real vector.

• scale_tril – a lower triangular matrix with positive diagonal.

domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._IndependentConstraint object>
log_abs_det_jacobian(x, y, intermediates=None)[source]
forward_shape(shape)[source]

Infers the shape of the forward computation, given the input shape. Defaults to preserving shape.

inverse_shape(shape)[source]

Infers the shapes of the inverse computation, given the output shape. Defaults to preserving shape.

### LowerCholeskyTransform¶

class LowerCholeskyTransform[source]

Transform a real vector to a lower triangular cholesky factor, where the strictly lower triangular submatrix is unconstrained and the diagonal is parameterized with an exponential transform.

domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._LowerCholesky object>
log_abs_det_jacobian(x, y, intermediates=None)[source]
forward_shape(shape)[source]

Infers the shape of the forward computation, given the input shape. Defaults to preserving shape.

inverse_shape(shape)[source]

Infers the shapes of the inverse computation, given the output shape. Defaults to preserving shape.

### OrderedTransform¶

class OrderedTransform[source]

Transform a real vector to an ordered vector.

References:

1. Stan Reference Manual v2.20, section 10.6, Stan Development Team

domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._OrderedVector object>
log_abs_det_jacobian(x, y, intermediates=None)[source]

### PermuteTransform¶

class PermuteTransform(permutation)[source]
domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._IndependentConstraint object>
log_abs_det_jacobian(x, y, intermediates=None)[source]

### PowerTransform¶

class PowerTransform(exponent)[source]
domain = <numpyro.distributions.constraints._GreaterThan object>
codomain = <numpyro.distributions.constraints._GreaterThan object>
log_abs_det_jacobian(x, y, intermediates=None)[source]
forward_shape(shape)[source]

Infers the shape of the forward computation, given the input shape. Defaults to preserving shape.

inverse_shape(shape)[source]

Infers the shapes of the inverse computation, given the output shape. Defaults to preserving shape.

### ScaledUnitLowerCholeskyTransform¶

class ScaledUnitLowerCholeskyTransform[source]

Like LowerCholeskyTransform this Transform transforms a real vector to a lower triangular cholesky factor. However it does so via a decomposition

$$y = loc + unit\_scale\_tril\ @\ scale\_diag\ @\ x$$.

where $$unit\_scale\_tril$$ has ones along the diagonal and $$scale\_diag$$ is a diagonal matrix with all positive entries that is parameterized with a softplus transform.

domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._ScaledUnitLowerCholesky object>
log_abs_det_jacobian(x, y, intermediates=None)[source]

### SigmoidTransform¶

class SigmoidTransform[source]
codomain = <numpyro.distributions.constraints._Interval object>
log_abs_det_jacobian(x, y, intermediates=None)[source]

### SimplexToOrderedTransform¶

class SimplexToOrderedTransform(anchor_point=0.0)[source]

Transform a simplex into an ordered vector (via difference in Logistic CDF between cutpoints) Used in [1] to induce a prior on latent cutpoints via transforming ordered category probabilities.

Parameters

anchor_point – Anchor point is a nuisance parameter to improve the identifiability of the transform. For simplicity, we assume it is a scalar value, but it is broadcastable x.shape[:-1]. For more details please refer to Section 2.2 in [1]

References:

1. Ordinal Regression Case Study, section 2.2, M. Betancourt, https://betanalpha.github.io/assets/case_studies/ordinal_regression.html

domain = <numpyro.distributions.constraints._Simplex object>
codomain = <numpyro.distributions.constraints._OrderedVector object>
log_abs_det_jacobian(x, y, intermediates=None)[source]

### SoftplusLowerCholeskyTransform¶

class SoftplusLowerCholeskyTransform[source]

Transform from unconstrained vector to lower-triangular matrices with nonnegative diagonal entries. This is useful for parameterizing positive definite matrices in terms of their Cholesky factorization.

domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._SoftplusLowerCholesky object>
log_abs_det_jacobian(x, y, intermediates=None)[source]
forward_shape(shape)[source]

Infers the shape of the forward computation, given the input shape. Defaults to preserving shape.

inverse_shape(shape)[source]

Infers the shapes of the inverse computation, given the output shape. Defaults to preserving shape.

### SoftplusTransform¶

class SoftplusTransform[source]

Transform from unconstrained space to positive domain via softplus $$y = \log(1 + \exp(x))$$. The inverse is computed as $$x = \log(\exp(y) - 1)$$.

domain = <numpyro.distributions.constraints._Real object>
codomain = <numpyro.distributions.constraints._SoftplusPositive object>
log_abs_det_jacobian(x, y, intermediates=None)[source]

### StickBreakingTransform¶

class StickBreakingTransform[source]
domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._Simplex object>
log_abs_det_jacobian(x, y, intermediates=None)[source]
forward_shape(shape)[source]

Infers the shape of the forward computation, given the input shape. Defaults to preserving shape.

inverse_shape(shape)[source]

Infers the shapes of the inverse computation, given the output shape. Defaults to preserving shape.

## Flows¶

### InverseAutoregressiveTransform¶

class InverseAutoregressiveTransform(autoregressive_nn, log_scale_min_clip=- 5.0, log_scale_max_clip=3.0)[source]

An implementation of Inverse Autoregressive Flow, using 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$$.

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

domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._IndependentConstraint object>
call_with_intermediates(x)[source]
log_abs_det_jacobian(x, y, intermediates=None)[source]

Calculates the elementwise determinant of the log jacobian.

Parameters

### BlockNeuralAutoregressiveTransform¶

class BlockNeuralAutoregressiveTransform(bn_arn)[source]

An implementation of Block Neural Autoregressive flow.

References

1. Block Neural Autoregressive Flow, Nicola De Cao, Ivan Titov, Wilker Aziz

domain = <numpyro.distributions.constraints._IndependentConstraint object>
codomain = <numpyro.distributions.constraints._IndependentConstraint object>
call_with_intermediates(x)[source]
log_abs_det_jacobian(x, y, intermediates=None)[source]

Calculates the elementwise determinant of the log jacobian.

Parameters