tfp.substrates.jax.distributions.JointDistributionNamedAutoBatched

Note: This API is new and only available in tfp-nightly.

Joint distribution parameterized by named distribution-making functions.

Inherits From: JointDistributionNamed

View aliases

Main aliases

tfp.experimental.substrates.jax.distributions.JointDistributionNamedAutoBatched

tfp.substrates.jax.distributions.JointDistributionNamedAutoBatched(
    model, batch_ndims=0, use_vectorized_map=True, validate_args=False, name=None
)

This class provides automatic vectorization and alternative semantics for tfd.JointDistributionNamed, 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.JointDistributionNamed model:

model = tfd.JointDistributionNamed({
  'x': tfd.Normal(0., tf.ones([3])),
  'y': tfd.Normal(0., 1.),
  'z': lambda x, y: tfd.Normal(x[..., :2] + y[..., tf.newaxis], 1.)
})

can be written in auto-vectorized form as

model = tfd.JointDistributionNamedAutoBatched({
  'x': tfd.Normal(0., tf.ones([3])),
  'y': tfd.Normal(0., 1.),
  'z': lambda x, y: tfd.Normal(x[:2] + y, 1.)
})

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

Note: auto-vectorization is still experimental and some TensorFlow ops may be unsupported. It can be disabled by setting use_vectorized_map=False.

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 JointDistributionNamedAutoBatched is the dictionary of tensors produced by .sample(); thus, the event_shape is the dictionary containing 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

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.JointDistributionNamedAutoBatched(dict(
    e=             tfd.Exponential(rate=[100, 120]),
    g=lambda    e: tfd.Gamma(concentration=e[0], rate=e[1]),
    n=             tfd.Normal(loc=0, scale=2.),
    m=lambda n, g: tfd.Normal(loc=n, scale=g),
    x=lambda    m: tfd.Sample(tfd.Bernoulli(logits=m), 12),
))

Notice the 1:1 correspondence between "math" and "code". In a standard JointDistributionNamed, we would have wrapped the first variable as e = tfd.Independent(tfd.Exponential(rate=[100, 120]), reinterpreted_batch_ndims=1) to specify that log_prob of the Exponential should be a scalar, summing over both dimensions. We would also have had to extend indices as tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]) to account for possible batch dimensions. Both of these behaviors are implicit in JointDistributionNamedAutoBatched.

Args
model	A generator that yields a sequence of tfd.Distribution-like instances.
batch_ndims	int Tensor number of batch dimensions. The batch_shapes 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.
name	The name for ops managed by the distribution. Default value: None (i.e., JointDistributionNamed).

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
dtype	The DType of Tensors handled by this Distribution.
event_shape
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

Methods

batch_shape_tensor

View source

batch_shape_tensor(
    sample_shape=(), name='batch_shape_tensor'
)

cdf

View source

cdf(
    value, name='cdf', **kwargs
)

Cumulative distribution function.

Given random variable X, the cumulative distribution function cdf is:

cdf(x) := P[X <= x]

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.

Returns
cdf	a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

copy

View source

copy(
    **override_parameters_kwargs
)

Creates a deep copy of the distribution.

Note: the copy distribution may continue to depend on the original initialization arguments.

Args
**override_parameters_kwargs	String/value dictionary of initialization arguments to override with new values.

Returns
distribution	A new instance of type(self) initialized from the union of self.parameters and override_parameters_kwargs, i.e., dict(self.parameters, **override_parameters_kwargs).

covariance

View source

covariance(
    name='covariance', **kwargs
)

Covariance.

Covariance is (possibly) defined only for non-scalar-event distributions.

For example, for a length-k, vector-valued distribution, it is calculated as,

Cov[i, j] = Covariance(X_i, X_j) = E[(X_i - E[X_i]) (X_j - E[X_j])]

where Cov is a (batch of) k x k matrix, 0 <= (i, j) < k, and E denotes expectation.

Alternatively, for non-vector, multivariate distributions (e.g., matrix-valued, Wishart), Covariance shall return a (batch of) matrices under some vectorization of the events, i.e.,

Cov[i, j] = Covariance(Vec(X)_i, Vec(X)_j) = [as above]

where Cov is a (batch of) k' x k' matrices, 0 <= (i, j) < k' = reduce_prod(event_shape), and Vec is some function mapping indices of this distribution's event dimensions to indices of a length-k' vector.

Args
name	Python str prepended to names of ops created by this function.
**kwargs	Named arguments forwarded to subclass implementation.

Returns
covariance	Floating-point Tensor with shape [B1, ..., Bn, k', k'] where the first n dimensions are batch coordinates and k' = reduce_prod(self.event_shape).

cross_entropy

View source

cross_entropy(
    other, name='cross_entropy'
)

Computes the (Shannon) cross entropy.

Denote this distribution (self) by P and the other distribution by Q. Assuming P, Q are absolutely continuous with respect to one another and permit densities p(x) dr(x) and q(x) dr(x), (Shannon) cross entropy is defined as:

H[P, Q] = E_p[-log q(X)] = -int_F p(x) log q(x) dr(x)

where F denotes the support of the random variable X ~ P.

other types with built-in registrations: JointDistributionNamed, JointDistributionNamedAutoBatched, JointDistributionSequential, JointDistributionSequentialAutoBatched

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

Returns
cross_entropy	self.dtype Tensor with shape [B1, ..., Bn] representing n different calculations of (Shannon) cross entropy.

entropy

View source

entropy(
    name='entropy', **kwargs
)

Shannon entropy in nats.

Additional documentation from JointDistributionSequential:

Shannon entropy in nats.

event_shape_tensor

View source

event_shape_tensor(
    sample_shape=(), name='event_shape_tensor'
)

Shape of a single sample from a single batch.

is_scalar_batch

View source

is_scalar_batch(
    name='is_scalar_batch'
)

Indicates that batch_shape == [].

Args
name	Python str prepended to names of ops created by this function.

Returns
is_scalar_batch	bool scalar Tensor.

is_scalar_event

View source

is_scalar_event(
    name='is_scalar_event'
)

Indicates that event_shape == [].

Args
name	Python str prepended to names of ops created by this function.

Returns
is_scalar_event	bool scalar Tensor for each distribution in model.

kl_divergence

View source

kl_divergence(
    other, name='kl_divergence'
)

Computes the Kullback--Leibler divergence.

Denote this distribution (self) by p and the other distribution by q. Assuming p, q are absolutely continuous with respect to reference measure r, the KL divergence is defined as:

KL[p, q] = E_p[log(p(X)/q(X))]
         = -int_F p(x) log q(x) dr(x) + int_F p(x) log p(x) dr(x)
         = H[p, q] - H[p]

where F denotes the support of the random variable X ~ p, H[., .] denotes (Shannon) cross entropy, and H[.] denotes (Shannon) entropy.

other types with built-in registrations: JointDistributionNamed, JointDistributionNamedAutoBatched, JointDistributionSequential, JointDistributionSequentialAutoBatched

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

Returns
kl_divergence	self.dtype Tensor with shape [B1, ..., Bn] representing n different calculations of the Kullback-Leibler divergence.

log_cdf

View source

log_cdf(
    value, name='log_cdf', **kwargs
)

Log cumulative distribution function.

Given random variable X, the cumulative distribution function cdf is:

log_cdf(x) := Log[ P[X <= x] ]

Often, a numerical approximation can be used for log_cdf(x) that yields a more accurate answer than simply taking the logarithm of the cdf when x << -1.

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.

Returns
logcdf	a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

log_prob

View source

log_prob(
    *args, **kwargs
)

Log probability density/mass function.

The measure methods of `JointDistribution` (`log_prob`, `prob`, etc.)
can be called either by passing a single structure of tensors or by using
named args for each part of the joint distribution state. For example,

jd = tfd.JointDistributionSequential([
    tfd.Normal(0., 1.),
    lambda z: tfd.Normal(z, 1.)
], validate_args=True)
jd.dtype
# ==> [tf.float32, tf.float32]
z, x = sample = jd.sample()
# The following calling styles are all permissable and produce the exactly
# the same output.
assert (jd.log_prob(sample) ==
        jd.log_prob(value=sample) ==
        jd.log_prob(z, x) ==
        jd.log_prob(z=z, x=x) ==
        jd.log_prob(z, x=x))

# These calling possibilities also imply that one can also use `*`
# expansion, if `sample` is a sequence:
jd.log_prob(*sample)
# and similarly, if `sample` is a map, one can use `**` expansion:
jd.log_prob(**sample)

`JointDistribution` component distributions names are resolved via
`jd._flat_resolve_names()`, which is implemented by each `JointDistribution`
subclass (see subclass documentation for details). Generally, for components
where a name was provided---
either explicitly as the `name` argument to a distribution or as a key in a
dict-valued JointDistribution, or implicitly, e.g., by the argument name of
a `JointDistributionSequential` distribution-making function---the provided
name will be used. Otherwise the component will receive a dummy name; these
may change without warning and should not be relied upon.

Note: not all `JointDistribution` subclasses support all calling styles;
for example, `JointDistributionNamed` does not support positional arguments
(aka "unnamed arguments") unless the provided model specifies an ordering of
variables (i.e., is an `collections.OrderedDict` or `collections.namedtuple`
rather than a plain `dict`).

Note: care is taken to resolve any potential ambiguity---this is generally
possible by inspecting the structure of the provided argument and "aligning"
it to the joint distribution output structure (defined by `jd.dtype`). For
example,

trivial_jd = tfd.JointDistributionSequential([tfd.Exponential(1.)])
trivial_jd.dtype  # => [tf.float32]
trivial_jd.log_prob([4.])
# ==> Tensor with shape `[]`.
lp = trivial_jd.log_prob(4.)
# ==> Tensor with shape `[]`.

Notice that in the first call, `[4.]` is interpreted as a list of one
scalar while in the second call the input is a scalar. Hence both inputs
result in identical scalar outputs. If we wanted to pass an explicit
vector to the `Exponential` component---creating a vector-shaped batch
of `log_prob`s---we could instead write
`trivial_jd.log_prob(np.array([4]))`.

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.

Returns
log_prob	a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

log_prob_parts

View source

log_prob_parts(
    value, name='log_prob_parts'
)

Log probability density/mass function.

Args
value	list of Tensors 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".

Returns
log_prob_parts	a tuple of Tensors representing the log_prob for each distribution_fn evaluated at each corresponding value.

log_survival_function

View source

log_survival_function(
    value, name='log_survival_function', **kwargs
)

Log survival function.

Given random variable X, the survival function is defined:

log_survival_function(x) = Log[ P[X > x] ]
                         = Log[ 1 - P[X <= x] ]
                         = Log[ 1 - cdf(x) ]

Typically, different numerical approximations can be used for the log survival function, which are more accurate than 1 - cdf(x) when x >> 1.

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.

Returns
Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

mean

View source

mean(
    name='mean', **kwargs
)

Mean.

mode

View source

mode(
    name='mode', **kwargs
)

Mode.

param_shapes

View source

@classmethod
param_shapes(
    sample_shape, name='DistributionParamShapes'
)

Shapes of parameters given the desired shape of a call to sample().

This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample().

Subclasses should override class method _param_shapes.

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

Returns
dict of parameter name to Tensor shapes.

param_static_shapes

View source

@classmethod
param_static_shapes(
    sample_shape
)

param_shapes with static (i.e. TensorShape) shapes.

This is a class method that describes what key/value arguments are required to instantiate the given Distribution so that a particular shape is returned for that instance's call to sample(). Assumes that the sample's shape is known statically.

Subclasses should override class method _param_shapes to return constant-valued tensors when constant values are fed.

Args
sample_shape	TensorShape or python list/tuple. Desired shape of a call to sample().

Returns
dict of parameter name to TensorShape.

Raises
ValueError	if sample_shape is a TensorShape and is not fully defined.

prob

View source

prob(
    *args, **kwargs
)

Probability density/mass function.

The measure methods of `JointDistribution` (`log_prob`, `prob`, etc.)
can be called either by passing a single structure of tensors or by using
named args for each part of the joint distribution state. For example,

jd = tfd.JointDistributionSequential([
    tfd.Normal(0., 1.),
    lambda z: tfd.Normal(z, 1.)
], validate_args=True)
jd.dtype
# ==> [tf.float32, tf.float32]
z, x = sample = jd.sample()
# The following calling styles are all permissable and produce the exactly
# the same output.
assert (jd.prob(sample) ==
        jd.prob(value=sample) ==
        jd.prob(z, x) ==
        jd.prob(z=z, x=x) ==
        jd.prob(z, x=x))

# These calling possibilities also imply that one can also use `*`
# expansion, if `sample` is a sequence:
jd.prob(*sample)
# and similarly, if `sample` is a map, one can use `**` expansion:
jd.prob(**sample)

`JointDistribution` component distributions names are resolved via
`jd._flat_resolve_names()`, which is implemented by each `JointDistribution`
subclass (see subclass documentation for details). Generally, for components
where a name was provided---
either explicitly as the `name` argument to a distribution or as a key in a
dict-valued JointDistribution, or implicitly, e.g., by the argument name of
a `JointDistributionSequential` distribution-making function---the provided
name will be used. Otherwise the component will receive a dummy name; these
may change without warning and should not be relied upon.

Note: not all `JointDistribution` subclasses support all calling styles;
for example, `JointDistributionNamed` does not support positional arguments
(aka "unnamed arguments") unless the provided model specifies an ordering of
variables (i.e., is an `collections.OrderedDict` or `collections.namedtuple`
rather than a plain `dict`).

Note: care is taken to resolve any potential ambiguity---this is generally
possible by inspecting the structure of the provided argument and "aligning"
it to the joint distribution output structure (defined by `jd.dtype`). For
example,

trivial_jd = tfd.JointDistributionSequential([tfd.Exponential(1.)])
trivial_jd.dtype  # => [tf.float32]
trivial_jd.prob([4.])
# ==> Tensor with shape `[]`.
prob = trivial_jd.prob(4.)
# ==> Tensor with shape `[]`.

Notice that in the first call, `[4.]` is interpreted as a list of one
scalar while in the second call the input is a scalar. Hence both inputs
result in identical scalar outputs. If we wanted to pass an explicit
vector to the `Exponential` component---creating a vector-shaped batch
of `prob`s---we could instead write
`trivial_jd.prob(np.array([4]))`.

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.

Returns
prob	a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

prob_parts

View source

prob_parts(
    value, name='prob_parts'
)

Log probability density/mass function.

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

Returns
prob_parts	a tuple of Tensors representing the prob for each distribution_fn evaluated at each corresponding value.

quantile

View source

quantile(
    value, name='quantile', **kwargs
)

Quantile function. Aka 'inverse cdf' or 'percent point function'.

Given random variable X and p in [0, 1], the quantile is:

quantile(p) := x such that P[X <= x] == p

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.

Returns
quantile	a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

resolve_graph

View source

resolve_graph(
    distribution_names=None, leaf_name='x'
)

Creates a tuple of tuples of dependencies.

This function is experimental. That said, we encourage its use and ask that you report problems to [e-mail].

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

Returns
graph	tuple of (str tuple) pairs representing the name of each distribution (maker) and the names of its dependencies.

Example

d = tfd.JointDistributionSequential([
                 tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
    lambda    e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
                 tfd.Normal(loc=0, scale=2.),
    lambda n, g: tfd.Normal(loc=n, scale=g),
])
d.resolve_graph()
# ==> (
#       ('e', ()),
#       ('g', ('e',)),
#       ('n', ()),
#       ('x', ('n', 'g')),
#     )

sample

View source

sample(
    sample_shape=(), seed=None, name='sample', **kwargs
)

Generate samples of the specified shape.

Note that a call to sample() without arguments will generate a single sample.

Args
sample_shape	0D or 1D int32 Tensor. Shape of the generated samples.
seed	Python integer or tfp.util.SeedStream instance, for seeding PRNG.
name	name to give to the op.
**kwargs	Named arguments forwarded to subclass implementation.

Returns
samples	a Tensor with prepended dimensions sample_shape.

sample_distributions

View source

sample_distributions(
    sample_shape=(), seed=None, value=None, name='sample_distributions'
)

stddev

View source

stddev(
    name='stddev', **kwargs
)

Standard deviation.

Standard deviation is defined as,

stddev = E[(X - E[X])**2]**0.5

where X is the random variable associated with this distribution, E denotes expectation, and stddev.shape = batch_shape + event_shape.

Args
name	Python str prepended to names of ops created by this function.
**kwargs	Named arguments forwarded to subclass implementation.

Returns
stddev	Floating-point Tensor with shape identical to batch_shape + event_shape, i.e., the same shape as self.mean().

survival_function

View source

survival_function(
    value, name='survival_function', **kwargs
)

Survival function.

Given random variable X, the survival function is defined:

survival_function(x) = P[X > x]
                     = 1 - P[X <= x]
                     = 1 - cdf(x).

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.

Returns
Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

variance

View source

variance(
    name='variance', **kwargs
)

Variance.

Variance is defined as,

Var = E[(X - E[X])**2]

where X is the random variable associated with this distribution, E denotes expectation, and Var.shape = batch_shape + event_shape.

Args
name	Python str prepended to names of ops created by this function.
**kwargs	Named arguments forwarded to subclass implementation.

Returns
variance	Floating-point Tensor with shape identical to batch_shape + event_shape, i.e., the same shape as self.mean().

__getitem__

View source

__getitem__(
    slices
)

Slices the batch axes of this distribution, returning a new instance.

b = tfd.Bernoulli(logits=tf.zeros([3, 5, 7, 9]))
b.batch_shape  # => [3, 5, 7, 9]
b2 = b[:, tf.newaxis, ..., -2:, 1::2]
b2.batch_shape  # => [3, 1, 5, 2, 4]

x = tf.random.stateless_normal([5, 3, 2, 2])
cov = tf.matmul(x, x, transpose_b=True)
chol = tf.linalg.cholesky(cov)
loc = tf.random.stateless_normal([4, 1, 3, 1])
mvn = tfd.MultivariateNormalTriL(loc, chol)
mvn.batch_shape  # => [4, 5, 3]
mvn.event_shape  # => [2]
mvn2 = mvn[:, 3:, ..., ::-1, tf.newaxis]
mvn2.batch_shape  # => [4, 2, 3, 1]
mvn2.event_shape  # => [2]

Args
slices	slices from the [] operator

Returns
dist	A new tfd.Distribution instance with sliced parameters.

__iter__

View source

__iter__()