tfp.distributions.JointDistributionCoroutineAutoBatched

Joint distribution parameterized by a distribution-making generator.

Inherits From: JointDistributionCoroutine, JointDistribution, AutoCompositeTensorDistribution, Distribution, AutoCompositeTensor

tfp.distributions.JointDistributionCoroutineAutoBatched(
    model,
    sample_dtype=None,
    batch_ndims=0,
    use_vectorized_map=True,
    validate_args=False,
    experimental_use_kahan_sum=False,
    name=None
)

Used in the notebooks

Used in the tutorials
Probabilistic PCA Approximate inference for STS models with non-Gaussian observations

This class provides automatic vectorization and alternative semantics for tfd.JointDistributionCoroutine, which in many cases allows for simplifications in the model specification.

Automatic vectorization

Auto-vectorized variants of JointDistribution allow the user to avoid explicitly annotating a model's vectorization semantics. 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 using tf.vectorized_map as required. In many cases this allows for significant simplications. For example, the following manually-vectorized tfd.JointDistributionCoroutine model:

def model_fn():
  x = yield tfd.JointDistributionCoroutine.Root(
    tfd.Normal(0., tf.ones([3])))
  y = yield tfd.JointDistributionCoroutine.Root(
    tfd.Normal(0., 1.))
  z = yield tfd.Normal(x[..., :2] + y[..., tf.newaxis], 1.)

can be written in auto-vectorized form as

```python
def model_fn():
  x = yield tfd.Normal(0., tf.ones([3]))
  y = yield tfd.Normal(0., 1.)
  z = yield tfd.Normal(x[:2] + y, 1.)

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

Alternative batch semantics

This class also provides alternative semantics for specifying a batch of independent (non-identical) joint distributions.

Instead of simply summing the log_probs of component distributions (which may have different shapes), it first reduces the component log_probs to ensure that jd.log_prob(jd.sample()) always returns a scalar, unless batch_ndims is explicitly set to a nonzero value (in which case the result will have the corresponding tensor rank).

The essential changes are:

An event of JointDistributionCoroutineAutoBatched is the list of tensors produced by .sample(); thus, the event_shape is the list of the shapes of sampled tensors. These combine both the event and batch dimensions of the component distributions. By contrast, the event shape of a base JointDistributions does not include batch dimensions of component distributions.
The batch_shape is a global property of the entire model, rather than a per-component property as in base JointDistributions. The global batch shape must be a prefix of the batch shapes of each component; the length of this prefix is specified by an optional argument batch_ndims. If batch_ndims is not specified, the model has batch shape [].

Examples

A hierarchical model of Poisson log-rates, written using tfd.JointDistributionCoroutineAutoBatched:

tfd = tfp.distributions
def model():
  global_log_rate = yield tfd.Normal(loc=0., scale=1.)
  local_log_rates = yield tfd.Normal(loc=0., scale=tf.ones([20]))
  observed_counts = yield tfd.Poisson(
    rate=tf.exp(global_log_rate + local_log_rates))
joint = tfd.JointDistributionCoroutineAutoBatched(model)

print(joint.event_shape)
# ==> [[], [20], [20]]
print(joint.batch_shape)
# ==> []
xs = joint.sample()
print([x.shape for x in xs])
# ==> [[], [20], [20]]
lp = joint.log_prob(xs)
print(lp.shape)
# ==> []

Note that the component distributions of this model would, by themselves, return batches of log-densities (because they are constructed with batch shape); the joint model implicitly sums over these to compute the single joint log-density.

ds, xs = joint.sample_distributions()
print([d.event_shape for d in ds])
# ==> [[], [], []] != model.event_shape
print([d.batch_shape for d in ds])
# ==> [[], [20], [20]] != model.batch_shape
print([d.log_prob(x).shape for (d, x) in zip(ds, xs)])
# ==> [[], [20], [20]]

The behavior of JointDistributionCoroutineAutoBatched is (assuming that batch_ndims is not specified) equivalent to adding tfp.distributions.Independent wrappers to reinterpret all batch dimensions in a JointDistributionCoroutine model. That is, the model above would be equivalently written using JointDistributionCoroutine as:

def model_jdc():
  global_log_rate = yield Root(tfd.Normal(0., 1.))
  local_log_rates = yield Root(tfd.Independent(
    tfd.Normal(0., tf.ones([20])), reinterpreted_batch_ndims=1))
  observed_counts = yield Root(tfd.Independent(
    tfd.Poisson(tf.exp(global_log_rate + local_log_rates)),
    reinterpreted_batch_ndims=1))
joint_jdc = tfd.JointDistributionCoroutine(model_jdc)

To define a batch of joint distributions (independent, but not identical, joint distributions from the same family) using JointDistributionCoroutineAutoBatched, any batch dimensions must be a shared prefix of the batch dimensions for all components. The batch_ndims argument determines the size of the prefix to consider. For example, consider a simple joint model with two scalar normal random variables, where the second variable's mean is given by the first variable. We can write a batch of five such models as:

def model():
  x = yield tfd.Normal(0., scale=tf.ones([5]))
  y = yield tfd.Normal(x, scale=[3., 2., 5., 1., 6.])
batch_joint = tfd.JointDistributionCoroutineAutoBatched(model, batch_ndims=1)

print(batch_joint.event_shape)
# ==> [[], []]
print(batch_joint.batch_shape)
# ==> [5]
print(batch_joint.log_prob(batch_joint.sample()).shape)
# ==> [5]

Note that if we had not passed batch_ndims, this would be interpreted as a single model over vector-valued random variables (whose components happen to be independent):

alternate_joint = tfd.JointDistributionCoroutineAutoBatched(model)
print(alternate_joint.event_shape)
# ==> [[5], [5]]
print(alternate_joint.batch_shape)
# ==> []
print(alternate_joint.log_prob(batch_joint.sample()).shape)
# ==> []

For improved readability of sampled values, the yielded distributions can also be named:

tfd = tfp.distributions
def model():
  global_log_rate = yield tfd.Normal(
    loc=0., scale=1., name='global_log_rate')
  local_log_rates = yield tfd.Normal(
    loc=0., scale=tf.ones([20]), name='local_log_rates')
  observed_counts = yield tfd.Poisson(
    rate=tf.exp(global_log_rate + local_log_rates), name='observed_counts')
joint = tfd.JointDistributionCoroutineAutoBatched(model)

print(joint.event_shape)
# ==> StructTuple(global_log_rate=[], local_log_rates=[20],
#      observed_counts=[20])
print(joint.batch_shape)
# ==> []
xs = joint.sample()
print(['{}: {}'.format(k, x.shape) for k, x in xs._asdict().items()])
# ==> global_log_scale: []
#     local_log_rates: [20]
#     observed_counts: [20]
lp = joint.log_prob(xs)
print(lp.shape)
# ==> []

# Passing via `kwargs` also works.
lp = joint.log_prob(**xs._asdict())
# Or:
lp = joint.log_prob(
    global_log_scale=...,
    local_log_rates=...,
    observed_counts=...,
)

If any of the yielded distributions are not explicitly named, they will automatically be given a name of the form var# where # is the index of the associated distribution. E.g. the first yielded distribution will have a default name of var0.

Args
`model`	A generator that yields a sequence of `tfd.Distribution`-like instances.
`sample_dtype`	Samples from this distribution will be structured like `tf.nest.pack_sequence_as(sample_dtype, list_)`. `sample_dtype` is only used for `tf.nest.pack_sequence_as` structuring of outputs, never casting (which is the responsibility of the component distributions). Default value: `None` (i.e. `namedtuple`).
`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: `0`.
`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: `True`.
`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`.
`name`	The name for ops managed by the distribution. Default value: `None` (i.e., `JointDistributionCoroutine`).

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`.
`name_scope`	Returns a `tf.name_scope` instance for this class.
`non_trainable_variables`	Sequence of non-trainable variables owned by this module and its submodules. Note: this method uses reflection to find variables on the current instance and submodules. For performance reasons you may wish to cache the result of calling this method if you don't expect the return value to change.
`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`.
`submodules`	Sequence of all sub-modules. Submodules are modules which are properties of this module, or found as properties of modules which are properties of this module (and so on). `a = tf.Module()` `b = tf.Module()` `c = tf.Module()` `a.b = b` `b.c = c` `list(a.submodules) == [b, c]` `True` `list(b.submodules) == [c]` `True` `list(c.submodules) == []` `True`
`trainable_variables`	Sequence of trainable variables owned by this module and its submodules. Note: this method uses reflection to find variables on the current instance and submodules. For performance reasons you may wish to cache the result of calling this method if you don't expect the return value to change.
`use_vectorized_map`
`validate_args`	Python `bool` indicating possibly expensive checks are enabled.
`variables`	Sequence of variables owned by this module and its submodules. Note: this method uses reflection to find variables on the current instance and submodules. For performance reasons you may wish to cache the result of calling this method if you don't expect the return value to change.

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
`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
`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.distributions.JointDistributionCoroutineAutoBatched

Used in the notebooks

Automatic vectorization

Alternative batch semantics

Examples

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

sample

kwargs:

sample_distributions

stddev

survival_function

unnormalized_log_prob

unnormalized_log_prob_parts

unnormalized_prob_parts

variance

with_name_scope

__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`

`sample`

`kwargs`:

`sample_distributions`

`stddev`

`survival_function`

`unnormalized_log_prob`

`unnormalized_log_prob_parts`

`unnormalized_prob_parts`

`variance`

`with_name_scope`

`getitem`

`iter`