tfp.substrates.jax.distributions.JointDistributionSequential

Joint distribution parameterized by distribution-making functions.

Inherits From: JointDistribution, Distribution

View aliases

Main aliases

tfp.experimental.substrates.jax.distributions.JointDistributionSequential

tfp.substrates.jax.distributions.JointDistributionSequential(
    model,
    batch_ndims=None,
    use_vectorized_map=False,
    validate_args=False,
    experimental_use_kahan_sum=False,
    name=None
)

This distribution enables both sampling and joint probability computation from a single model specification.

A joint distribution is a collection of possibly interdependent distributions. Like tf.keras.Sequential, the JointDistributionSequential can be specified via a list of functions (each responsible for making a tfp.distributions.Distribution-like instance). Unlike tf.keras.Sequential, each function can depend on the output of all previous elements rather than only the immediately previous.

Mathematical Details

The JointDistributionSequential implements the chain rule of probability. That is, the probability function of a length-d vector x is,

p(x) = prod{ p(x[i] | x[:i]) : i = 0, ..., (d - 1) }

The JointDistributionSequential is parameterized by a list comprised of either:

tfp.distributions.Distribution-like instances or,
callables which return a tfp.distributions.Distribution-like instance.

Each list element implements the i-th full conditional distribution, p(x[i] | x[:i]). The "conditioned on" elements are represented by the callable's required arguments. Directly providing a Distribution-like instance is a convenience and is semantically identical a zero argument callable.

Denote the i-th callables non-default arguments as args[i]. Since the callable is the conditional manifest, 0 <= len(args[i]) <= i - 1. When len(args[i]) < i - 1, the callable only depends on a subset of the previous distributions, specifically those at indexes: range(i - 1, i - 1 - num_args[i], -1). (See "Examples" and "Discussion" for why the order is reversed.)

Name resolution: The names ofJointDistributionSequentialcomponents are defined by explicitnamearguments passed to distributions (tfd.Normal(0., 1., name='x')) and/or by the argument names in distribution-making functions (lambda x: tfd.Normal(x., 1.)). Both approaches may be used in the same distribution, as long as they are consistent; referring to a single component by multiple names will raise aValueError`. Unnamed components will be assigned a dummy name.

Examples

Consider the following generative model:

e ~ Exponential(rate=[100,120])
g ~ Gamma(concentration=e[0], rate=e[1])
n ~ Normal(loc=0, scale=2.)
m ~ Normal(loc=n, scale=g)
for i = 1, ..., 12:
  x[i] ~ Bernoulli(logits=m)

We can code this as:

tfd = tfp.distributions
joint = tfd.JointDistributionSequential([
                 tfd.Exponential(rate=[100, 120]),           # e
    lambda    e: tfd.Gamma(concentration=e[0], rate=e[1]),    # g
                 tfd.Normal(loc=0, scale=2.),                 # n
    lambda n, g: tfd.Normal(loc=n, scale=g)                   # m
    lambda    m: tfd.Sample(tfd.Bernoulli(logits=m), 12)      # x
], batch_ndims=0, use_vectorized_map=True)

Notice the 1:1 correspondence between "math" and "code".

x = joint.sample()
# ==> A length-5 list of Tensors representing a draw/realization from each
#     distribution.
joint.log_prob(x)
# ==> A scalar `Tensor` representing the total log prob under all five
#     distributions.

joint.resolve_graph()
# ==> (('e', ()),
#      ('g', ('e',)),
#      ('n', ()),
#      ('m', ('n', 'g')),
#      ('x', ('m',)))

Discussion

JointDistributionSequential builds each distribution in list order; list items must be either a:

tfd.Distribution-like instance (e.g., e and n), or a
Python callable (e.g., g, m, x).

Regarding #1, an object is deemed "tfd.Distribution-like" if it has a sample, log_prob, and distribution properties, e.g., batch_shape, event_shape, dtype.

Regarding #2, in addition to using a function (or lambda), supplying a TFD "class" is also permissible, this also being a "Python callable." For example, instead of writing: lambda loc, scale: tfd.Normal(loc=loc, scale=scale) one could have simply written tfd.Normal.

Notice that directly providing a tfd.Distribution-like instance means there cannot exist a (dynamic) dependency on other distributions; it is "independent" both "computationally" and "statistically." The same is self-evidently true of zero-argument callables.

A distribution instance depends on other distribution instances through the distribution making function's required arguments. If the distribution maker has k required arguments then the JointDistributionSequential calls the maker with samples produced by the previous k distributions.

Vectorized sampling and model evaluation

When a joint distribution's sample method is called with a sample_shape (or the log_prob method is called on an input with multiple sample dimensions) the model must be equipped to handle additional batch dimensions. This may be done manually, or automatically by passing use_vectorized_map=True. Manual vectorization has historically been the default, but we now recommend that most users enable automatic vectorization unless they are affected by a specific issue; some known issues are listed below.

When using manually-vectorized joint distributions, each operation in the model must account for the possibility of batch dimensions in Distributions and their samples. By contrast, auto-vectorized models need only describe a single sample from the joint distribution; any batch evaluation is automated as required using tf.vectorized_map (vmap in JAX). In many cases this allows for significant simplications. For example, the following manually-vectorized tfd.JointDistributionSequential model:

model = tfd.JointDistributionSequential([
    tfd.Normal(0., tf.ones([3])),
    tfd.Normal(0., 1.),
    lambda y, x: tfd.Normal(x[..., :2] + y[..., tf.newaxis], 1.)
  ])

can be written in auto-vectorized form as

model = tfd.JointDistributionSequential([
    tfd.Normal(0., tf.ones([3])),
    tfd.Normal(0., 1.),
    lambda y, x: tfd.Normal(x[:2] + y, 1.)
  ],
  use_vectorized_map=True)

in which we were able to avoid explicitly accounting for batch dimensions when indexing and slicing computed quantities in the third line.

Known limitations of automatic vectorization:

A small fraction of TensorFlow ops are unsupported; models that use an unsupported op will raise an error and must be manually vectorized.
Sampling large batches may be slow under automatic vectorization because TensorFlow's stateless samplers are currently converted using a non-vectorized while_loop. This limitation applies only in TensorFlow; vectorized samplers in JAX should be approximately as fast as manually vectorized code.
Calling sample_distributions with nontrivial sample_shape will raise an error if the model contains any distributions that are not registered as CompositeTensors (TFP's basic distributions are usually fine, but support for wrapper distributions like tfd.Sample is a work in progress).

Batch semantics and (log-)densities

tl;dr: pass batch_ndims=0 unless you have a good reason not to.

Joint distributions now support 'auto-batching' semantics, in which the distribution's batch shape is derived by broadcasting the leftmost batch_ndims dimensions of its components' batch shapes. All remaining dimensions are considered to form a single 'event' of the joint distribution. If batch_ndims==0, then the joint distribution has batch shape [], and all component dimensions are treated as event shape. For example, the model

jd = tfd.JointDistributionSequential([
    tfd.Normal(0., tf.ones([3])),
    lambda x: tfd.Normal(x[..., tf.newaxis], tf.ones([3, 2]))
  ],
  batch_ndims=0)

creates a joint distribution with batch shape [] and event shape ([3], [3, 2]). The log-density of a sample always has shape batch_shape, so this guarantees that jd.log_prob(jd.sample()) will evaluate to a scalar value. We could alternately construct a joint distribution with batch shape [3] and event shape ([], [2]) by setting batch_ndims=1, in which case jd.log_prob(jd.sample()) would evaluate to a value of shape [3].

Setting batch_ndims=None recovers the 'classic' batch semantics (currently still the default for backwards-compatibility reasons), in which the joint distribution's log_prob is computed by naively summing log densities from the component distributions. Since these component densities have shapes equal to the batch shapes of the individual components, to avoid broadcasting errors it is usually necessary to construct the components with identical batch shapes. For example, the component distributions in the model above have batch shapes of [3] and [3, 2] respectively, which would raise an error if summed directly, but can be aligned by wrapping with tfd.Independent, as in this model:

jd = tfd.JointDistributionSequential([
    tfd.Normal(0., tf.ones([3])),
    lambda x: tfd.Independent(tfd.Normal(x[..., tf.newaxis], tf.ones([3, 2])),
                              reinterpreted_batch_ndims=1)
  ],
  batch_ndims=None)

Here the components both have batch shape [3], so jd.log_prob(jd.sample()) returns a value of shape [3], just as in the batch_ndims=1 case above. In fact, auto-batching semantics are equivalent to implicitly wrapping each component dist as tfd.Independent(dist, reinterpreted_batch_ndim=(dist.batch_shape.ndims - jd.batch_ndims)); the only vestigial difference is that under auto-batching semantics, the joint distribution has a single batch shape [3], while under the classic semantics the value of jd.batch_shape is a structure of the component batch shapes ([3], [3]). Such structured batch shapes will be deprecated in the future, since they are inconsistent with the definition of batch shapes used elsewhere in TFP.

References

[1] Dan Piponi, Dave Moore, and Joshua V. Dillon. Joint distributions for TensorFlow Probability. arXiv preprint arXiv:2001.11819_, 2020. https://arxiv.org/abs/2001.11819

Args
`model`	Python list of either tfd.Distribution instances and/or lambda functions which take the `k` previous distributions and returns a new tfd.Distribution instance.
`batch_ndims`	`int` `Tensor` number of batch dimensions. The `batch_shape`s of all component distributions must be such that the prefixes of length `batch_ndims` broadcast to a consistent joint batch shape. Default value: `None`.
`use_vectorized_map`	Python `bool`. Whether to use `tf.vectorized_map` to automatically vectorize evaluation of the model. This allows the model specification to focus on drawing a single sample, which is often simpler, but some ops may not be supported. Default value: `False`.
`validate_args`	Python `bool`. Whether to validate input with asserts. If `validate_args` is `False`, and the inputs are invalid, correct behavior is not guaranteed. Default value: `False`.
`experimental_use_kahan_sum`	Python `bool`. When `True`, we use Kahan summation to aggregate independent underlying log_prob values, which improves against the precision of a naive float32 sum. This can be noticeable in particular for large dimensions in float32. See CPU caveat on `tfp.math.reduce_kahan_sum`. This argument has no effect if `batch_ndims is None`. Default value: `False`.
`name`	The name for ops managed by the distribution. Default value: `None` (i.e., `"JointDistributionSequential"`).

Attributes
`allow_nan_stats`	Python `bool` describing behavior when a stat is undefined. Stats return +/- infinity when it makes sense. E.g., the variance of a Cauchy distribution is infinity. However, sometimes the statistic is undefined, e.g., if a distribution's pdf does not achieve a maximum within the support of the distribution, the mode is undefined. If the mean is undefined, then by definition the variance is undefined. E.g. the mean for Student's T for df = 1 is undefined (no clear way to say it is either + or - infinity), so the variance = E[(X - mean)**2] is also undefined.
`batch_ndims`
`batch_shape`	Shape of a single sample from a single event index as a `TensorShape`. May be partially defined or unknown. The batch dimensions are indexes into independent, non-identical parameterizations of this distribution.
`dtype`	The `DType` of `Tensor`s handled by this `Distribution`.
`event_shape`	Shape of a single sample from a single batch as a `TensorShape`. May be partially defined or unknown.
`experimental_shard_axis_names`	Indicates whether part distributions have active shard axis names.
`model`
`name`	Name prepended to all ops created by this `Distribution`.
`parameters`	Dictionary of parameters used to instantiate this `Distribution`.
`reparameterization_type`	Describes how samples from the distribution are reparameterized. Currently this is one of the static instances `tfd.FULLY_REPARAMETERIZED` or `tfd.NOT_REPARAMETERIZED`.
`trainable_variables`
`use_vectorized_map`
`validate_args`	Python `bool` indicating possibly expensive checks are enabled.
`variables`

Args
`value`	`float` or `double` `Tensor`.
`name`	Python `str` prepended to names of ops created by this function.
`**kwargs`	Named arguments forwarded to subclass implementation.

Args
`other`	`tfp.distributions.Distribution` instance.
`name`	Python `str` prepended to names of ops created by this function.

Args
`*args`	Passed to implementation `_default_event_space_bijector`.
`**kwargs`	Passed to implementation `_default_event_space_bijector`.

Args
`value`	a `Tensor` valid sample from this distribution family.
`sample_ndims`	Positive `int` Tensor number of leftmost dimensions of `value` that index i.i.d. samples. Default value: `1`.
`validate_args`	Python `bool`, default `False`. When `True`, distribution parameters are checked for validity despite possibly degrading runtime performance. When `False`, invalid inputs may silently render incorrect outputs. Default value: `False`.
`**init_kwargs`	Additional keyword arguments passed through to `cls.__init__`. These take precedence in case of collision with the fitted parameters; for example, `tfd.Normal.experimental_fit([1., 1.], scale=20.)` returns a Normal distribution with `scale=20.` rather than the maximum likelihood parameter `scale=0.`.

Args
`value`	`float` or `double` `Tensor`.
`backward_compat`	`bool` specifying whether to fall back to returning `FullSpace` as the tangent space, and representing R^n with the standard basis.
`**kwargs`	Named arguments forwarded to subclass implementation.

Returns
`log_prob`	a `Tensor` representing the log probability density, of shape `sample_shape(x) + self.batch_shape` with values of type `self.dtype`.
`tangent_space`	a `TangentSpace` object (by default `FullSpace`) representing the tangent space to the manifold at `value`.

Args
`*args`	Positional arguments: a value structure or component values (see above).
`**kwargs`	Keyword arguments: a value structure or component values (see above). May also include `name`, specifying a Python string name for ops generated by this method.

Args
`sample_shape`	integer `Tensor` desired shape of samples to draw. Default value: `()`.
`seed`	PRNG seed; see `tfp.random.sanitize_seed` for details. Default value: `None`.
`name`	name to give to the op. Default value: `'sample_and_log_prob'`.
`**kwargs`	Named arguments forwarded to subclass implementation.

Returns
`samples`	a `Tensor`, or structure of `Tensor`s, with prepended dimensions `sample_shape`.
`log_prob`	a `Tensor` of shape `sample_shape(x) + self.batch_shape` with values of type `self.dtype`.

Args
`*args`	Positional arguments: a `value` structure or component values (see above).
`**kwargs`	Keyword arguments: a `value` structure or component values (see above). May also include `name`, specifying a Python string name for ops generated by this method.

Args
`value`	`list` of `Tensor`s in `distribution_fn` order for which we compute the `log_prob_parts` and to parameterize other ("downstream") distributions.
`name`	name prepended to ops created by this function. Default value: `"log_prob_parts"`.

Args
`sample_shape`	`Tensor` or python list/tuple. Desired shape of a call to `sample()`.
`name`	name to prepend ops with.

Args
`dtype`	Optional float `dtype` to assume for continuous-valued parameters. Some constraining bijectors require advance knowledge of the dtype because certain constants (e.g., `tfb.Softplus.low`) must be instantiated with the same dtype as the values to be transformed.
`num_classes`	Optional `int` `Tensor` number of classes to assume when inferring the shape of parameters for categorical-like distributions. Otherwise ignored.

Args
`distribution_names`	`list` of `str` or `None` names corresponding to each of `model` elements. (`None`s are expanding into the appropriate `str`.)
`leaf_name`	`str` used when no maker depends on a particular `model` element.

Args
`sample_shape`	0D or 1D `int32` `Tensor`. Shape of the generated samples.
`seed`	PRNG seed; see `tfp.random.sanitize_seed` for details.
`name`	name to give to the op.
`**kwargs`	Named arguments forwarded to subclass implementation.

Returns
`distributions`	a `tuple` of `Distribution` instances for each of `distribution_fn`.
`samples`	a `tuple` of `Tensor`s with prepended dimensions `sample_shape` for each of `distribution_fn`.

tfp.substrates.jax.distributions.JointDistributionSequential

View aliases

Mathematical Details

Examples

We can code this as:

Discussion

Vectorized sampling and model evaluation

Batch semantics and (log-)densities

References

Args

Attributes

Child Classes

Methods

batch_shape_tensor

cdf

copy

covariance

cross_entropy

entropy

event_shape_tensor

experimental_default_event_space_bijector

experimental_fit

experimental_local_measure

experimental_pin

Examples:

experimental_sample_and_log_prob

is_scalar_batch

is_scalar_event

kl_divergence

log_cdf

log_prob

log_prob_parts

log_survival_function

mean

mode

param_shapes

param_static_shapes

parameter_properties

prob

prob_parts

quantile

resolve_graph

Example

sample

kwargs:

sample_distributions

stddev

survival_function

unnormalized_log_prob

unnormalized_log_prob_parts

unnormalized_prob_parts

variance

__getitem__

__iter__

`batch_shape_tensor`

`cdf`

`copy`

`covariance`

`cross_entropy`

`entropy`

`event_shape_tensor`

`experimental_default_event_space_bijector`

`experimental_fit`

`experimental_local_measure`

`experimental_pin`

`experimental_sample_and_log_prob`

`is_scalar_batch`

`is_scalar_event`

`kl_divergence`

`log_cdf`

`log_prob`

`log_prob_parts`

`log_survival_function`

`mean`

`mode`

`param_shapes`

`param_static_shapes`

`parameter_properties`

`prob`

`prob_parts`

`quantile`

`resolve_graph`

`sample`

`kwargs`:

`sample_distributions`

`stddev`

`survival_function`

`unnormalized_log_prob`

`unnormalized_log_prob_parts`

`unnormalized_prob_parts`

`variance`

`getitem`

`iter`