Base Distribution¶
Distribution¶
-
class
Distribution
(batch_shape=(), event_shape=(), validate_args=None)[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:
>>> d = dist.Dirichlet(np.ones((2, 3, 4))) >>> d.batch_shape (2, 3) >>> d.event_shape (4,)
-
arg_constraints
= {}¶
-
support
= None¶
-
reparametrized_params
= []¶
-
batch_shape
¶ Returns the shape over which the distribution parameters are batched.
Returns: batch shape of the distribution. Return type: tuple
-
event_shape
¶ Returns the shape of a single sample from the distribution without batching.
Returns: event shape of the distribution. 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 to be used for the distribution.
- size – the sample shape for the distribution.
Returns: a numpy.ndarray of shape sample_shape + batch_shape + event_shape
-
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: a numpy.ndarray with shape value.shape[:-self.event_shape]
-
mean
¶ Mean of the distribution.
-
variance
¶ Variance of the distribution.
TransformedDistribution¶
-
class
TransformedDistribution
(base_distribution, transforms, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
Returns a distribution instance obtained as a result of applying a sequence of transforms to a base distribution. For an example, see
LogNormal
andHalfNormal
.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
= {}¶
-
support
¶
-
is_reparametrized
¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
mean
¶
-
variance
¶
Continuous Distributions¶
Beta¶
-
class
Beta
(concentration1, concentration0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'concentration0': <numpyro.distributions.constraints._GreaterThan object>, 'concentration1': <numpyro.distributions.constraints._GreaterThan object>}¶
-
support
= <numpyro.distributions.constraints._Interval object>¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
Cauchy¶
-
class
Cauchy
(loc=0.0, scale=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
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]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
Chi2¶
-
class
Chi2
(df, validate_args=None)[source]¶ Bases:
numpyro.distributions.continuous.Gamma
-
arg_constraints
= {'df': <numpyro.distributions.constraints._GreaterThan object>}¶
-
Dirichlet¶
-
class
Dirichlet
(concentration, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'concentration': <numpyro.distributions.constraints._GreaterThan object>}¶
-
support
= <numpyro.distributions.constraints._Simplex object>¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
Exponential¶
-
class
Exponential
(rate=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
reparametrized_params
= ['rate']¶
-
arg_constraints
= {'rate': <numpyro.distributions.constraints._GreaterThan object>}¶
-
support
= <numpyro.distributions.constraints._GreaterThan object>¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
Gamma¶
-
class
Gamma
(concentration, rate=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'concentration': <numpyro.distributions.constraints._GreaterThan object>, 'rate': <numpyro.distributions.constraints._GreaterThan object>}¶
-
support
= <numpyro.distributions.constraints._GreaterThan object>¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
GaussianRandomWalk¶
-
class
GaussianRandomWalk
(scale=1.0, num_steps=1, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'num_steps': <numpyro.distributions.constraints._IntegerGreaterThan object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}¶
-
support
= <numpyro.distributions.constraints._Real object>¶
-
reparametrized_params
= ['scale']¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
HalfCauchy¶
-
class
HalfCauchy
(scale=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.TransformedDistribution
-
reparametrized_params
= ['scale']¶
-
arg_constraints
= {'scale': <numpyro.distributions.constraints._GreaterThan object>}¶
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
HalfNormal¶
-
class
HalfNormal
(scale=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.TransformedDistribution
-
reparametrized_params
= ['scale']¶
-
arg_constraints
= {'scale': <numpyro.distributions.constraints._GreaterThan object>}¶
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
LKJCholesky¶
-
class
LKJCholesky
(dimension, concentration=1.0, sample_method='onion', validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
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, whenconcentration == 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.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>}¶
-
support
= <numpyro.distributions.constraints._CorrCholesky object>¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
LogNormal¶
-
class
LogNormal
(loc=0.0, scale=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.TransformedDistribution
-
arg_constraints
= {'loc': <numpyro.distributions.constraints._Real object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}¶
-
reparametrized_params
= ['loc', 'scale']¶
-
Normal¶
-
class
Normal
(loc=0.0, scale=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
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]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
Pareto¶
-
class
Pareto
(alpha, scale=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.TransformedDistribution
-
arg_constraints
= {'alpha': <numpyro.distributions.constraints._GreaterThan object>, 'scale': <numpyro.distributions.constraints._GreaterThan object>}¶
-
support
¶
-
StudentT¶
-
class
StudentT
(df, loc=0.0, scale=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
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
= ['loc', 'scale']¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
TruncatedCauchy¶
-
class
TruncatedCauchy
(low=0.0, loc=0.0, scale=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
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']¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
support
¶
-
TruncatedNormal¶
-
class
TruncatedNormal
(low=0.0, loc=0.0, scale=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
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']¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
support
¶
-
Uniform¶
-
class
Uniform
(low=0.0, high=1.0, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'high': <numpyro.distributions.constraints._Dependent object>, 'low': <numpyro.distributions.constraints._Dependent object>}¶
-
reparametrized_params
= ['low', 'high']¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
support
¶
-
Discrete Distributions¶
BernoulliLogits¶
-
class
BernoulliLogits
(logits=None, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'logits': <numpyro.distributions.constraints._Real object>}¶
-
support
= <numpyro.distributions.constraints._Boolean object>¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
BernoulliProbs¶
-
class
BernoulliProbs
(probs, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'probs': <numpyro.distributions.constraints._Interval object>}¶
-
support
= <numpyro.distributions.constraints._Boolean object>¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
BinomialLogits¶
-
class
BinomialLogits
(logits, total_count=1, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'logits': <numpyro.distributions.constraints._Real object>, 'total_count': <numpyro.distributions.constraints._IntegerGreaterThan object>}¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
support
¶
-
BinomialProbs¶
-
class
BinomialProbs
(probs, total_count=1, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'probs': <numpyro.distributions.constraints._Interval object>, 'total_count': <numpyro.distributions.constraints._IntegerGreaterThan object>}¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
support
¶
-
CategoricalLogits¶
-
class
CategoricalLogits
(logits, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'logits': <numpyro.distributions.constraints._Real object>}¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
support
¶
-
CategoricalProbs¶
-
class
CategoricalProbs
(probs, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'probs': <numpyro.distributions.constraints._Simplex object>}¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
support
¶
-
MultinomialLogits¶
-
class
MultinomialLogits
(logits, total_count=1, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'logits': <numpyro.distributions.constraints._Real object>, 'total_count': <numpyro.distributions.constraints._IntegerGreaterThan object>}¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
support
¶
-
MultinomialProbs¶
-
class
MultinomialProbs
(probs, total_count=1, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'probs': <numpyro.distributions.constraints._Simplex object>, 'total_count': <numpyro.distributions.constraints._IntegerGreaterThan object>}¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-
support
¶
-
Poisson¶
-
class
Poisson
(rate, validate_args=None)[source]¶ Bases:
numpyro.distributions.distribution.Distribution
-
arg_constraints
= {'rate': <numpyro.distributions.constraints._GreaterThan object>}¶
-
support
= <numpyro.distributions.constraints._IntegerGreaterThan object>¶
-
sample
(key, sample_shape=())[source]¶ See
numpyro.distributions.distribution.Distribution.sample()
-
log_prob
(value)[source]¶ See
numpyro.distributions.distribution.Distribution.log_prob()
-