tfp.substrates.jax.distributions.TransformedDistribution

A Transformed Distribution.

Inherits From: Distribution

A TransformedDistribution models p(y) given a base distribution p(x), and a deterministic, invertible, differentiable transform, Y = g(X). The transform is typically an instance of the Bijector class and the base distribution is typically an instance of the Distribution class.

A Bijector is expected to implement the following functions:

  • forward,
  • inverse,
  • inverse_log_det_jacobian.

The semantics of these functions are outlined in the Bijector documentation.

We now describe how a TransformedDistribution alters the input/outputs of a Distribution associated with a random variable (rv) X.

Write cdf(Y=y) for an absolutely continuous cumulative distribution function of random variable Y; write the probability density function pdf(Y=y) := d^k / (dy_1,...,dy_k) cdf(Y=y) for its derivative wrt to Y evaluated at y. Assume that Y = g(X) where g is a deterministic diffeomorphism, i.e., a non-random, continuous, differentiable, and invertible function. Write the inverse of g as X = g^{-1}(Y) and (J o g)(x) for the Jacobian of g evaluated at x.

A TransformedDistribution implements the following operations:

  • sample Mathematically: Y = g(X) Programmatically: bijector.forward(distribution.sample(...))

  • log_prob Mathematically: `(log o pdf)(Y=y) = (log o pdf o g^{-1})(y)

                   + (log o abs o det o J o g^{-1})(y)`
    

    Programmatically: (distribution.log_prob(bijector.inverse(y)) + bijector.inverse_log_det_jacobian(y))

  • log_cdf Mathematically: (log o cdf)(Y=y) = (log o cdf o g^{-1})(y) Programmatically: distribution.log_cdf(bijector.inverse(x))

  • and similarly for: cdf, prob, log_survival_function, survival_function.

A simple example constructing a Log-Normal distribution from a Normal distribution:

tfd = tfp.distributions
tfb = tfp.bijectors
log_normal = tfd.TransformedDistribution(
  distribution=tfd.Normal(loc=0., scale=1.),
  bijector=tfb.Exp(),
  name='LogNormalTransformedDistribution')

A LogNormal made from callables:

tfd = tfp.distributions
tfb = tfp.bijectors
log_normal = tfd.TransformedDistribution(
  distribution=tfd.Normal(loc=0., scale=1.),
  bijector=tfb.Inline(
    forward_fn=tf.exp,
    inverse_fn=tf.log,
    inverse_log_det_jacobian_fn=(
      lambda y: -tf.reduce_sum(tf.log(y), axis=-1)),
  name='LogNormalTransformedDistribution')

Another example constructing a Normal from a StandardNormal:

tfd = tfp.distributions
tfb = tfp.bijectors
normal = tfd.TransformedDistribution(
  distribution=tfd.Normal(loc=0., scale=1.),
  bijector=tfb.Affine(
    shift=-1.,
    scale_identity_multiplier=2.)
  name='NormalTransformedDistribution')

A TransformedDistribution's batch_shape is the same as that of the base distribution, and its event_shape is the forward_event_shape of the bijector applied to the event_shape of the base distribution.

tfd.Sample or tfd.Independent may be used to add extra IID dimensions to the event_shape of the base distribution before the bijector operates on it. The following example demonstrates how to construct a multivariate Normal as a TransformedDistribution, by adding a rank-1 IID dimension to the event_shape of a standard Normal and applying tfb.ScaleMatvecTriL.

tfd = tfp.distributions
tfb = tfp.bijectors
# We will create two MVNs with batch_shape = event_shape = 2.
mean = [[-1., 0],      # batch:0
        [0., 1]]       # batch:1
chol_cov = [[[1., 0],
             [0, 1]],  # batch:0
            [[1, 0],
             [2, 2]]]  # batch:1
mvn1 = tfd.TransformedDistribution(
    distribution=tfd.Sample(
        tfd.Normal(loc=[0., 0], scale=1.),  # base_dist.batch_shape == [2]
        sample_shape=[2])                   # base_dist.event_shape == [2]
    bijector=tfb.Shift(shift=mean)(tfb.ScaleMatvecTriL(scale_tril=chol_cov)))
mvn2 = ds.MultivariateNormalTriL(loc=mean, scale_tril=chol_cov)
# mvn1.log_prob(x) == mvn2.log_prob(x)

<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2"><h2 class="add-link">Args</h2></th></tr>

<tr>
<td>
`distribution`
</td>
<td>
The base distribution instance to transform. Typically an
instance of `Distribution`.
</td>
</tr><tr>
<td>
`bijector`
</td>
<td>
The object responsible for calculating the transformation.
Typically an instance of `Bijector`.
</td>
</tr><tr>
<td>
`kwargs_split_fn`
</td>
<td>
Python `callable` which takes a kwargs `dict` and returns
a tuple of kwargs `dict`s for each of the `distribution` and `bijector`
parameters respectively.
Default value: `_default_kwargs_split_fn` (i.e.,
`lambda kwargs: (kwargs.get('distribution_kwargs', {}),
kwargs.get('bijector_kwargs', {}))`)
</td>
</tr><tr>
<td>
`validate_args`
</td>
<td>
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.
</td>
</tr><tr>
<td>
`parameters`
</td>
<td>
Locals dict captured by subclass constructor, to be used for
copy/slice re-instantiation operations.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
Python `str` name prefixed to Ops created by this class. Default:
`bijector.name + distribution.name`.
</td>
</tr>
</table>





<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2"><h2 class="add-link">Attributes</h2></th></tr>

<tr>
<td>
`allow_nan_stats`
</td>
<td>
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.
</td>
</tr><tr>
<td>
`batch_shape`
</td>
<td>
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.
</td>
</tr><tr>
<td>
`bijector`
</td>
<td>
Function transforming x => y.
</td>
</tr><tr>
<td>
`distribution`
</td>
<td>
Base distribution, p(x).
</td>
</tr><tr>
<td>
`dtype`
</td>
<td>
The `DType` of `Tensor`s handled by this `Distribution`.
</td>
</tr><tr>
<td>
`event_shape`
</td>
<td>
Shape of a single sample from a single batch as a `TensorShape`.

May be partially defined or unknown.
</td>
</tr><tr>
<td>
`name`
</td>
<td>
Name prepended to all ops created by this `Distribution`.
</td>
</tr><tr>
<td>
`parameters`
</td>
<td>
Dictionary of parameters used to instantiate this `Distribution`.
</td>
</tr><tr>
<td>
`reparameterization_type`
</td>
<td>
Describes how samples from the distribution are reparameterized.

Currently this is one of the static instances
`tfd.FULLY_REPARAMETERIZED` or `tfd.NOT_REPARAMETERIZED`.
</td>
</tr><tr>
<td>
`trainable_variables`
</td>
<td>

</td>
</tr><tr>
<td>
`validate_args`
</td>
<td>
Python `bool` indicating possibly expensive checks are enabled.
</td>
</tr><tr>
<td>
`variables`
</td>
<td>

</td>
</tr>
</table>



## Methods

<h3 id="batch_shape_tensor"><code>batch_shape_tensor</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/master/tensorflow_probability/substrates/jax/distributions/distribution.py">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>batch_shape_tensor(
    name='batch_shape_tensor'
)
</code></pre>

Shape of a single sample from a single event index as a 1-D `Tensor`.

The batch dimensions are indexes into independent, non-identical
parameterizations of this distribution.

<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Args</th></tr>

<tr>
<td>
`name`
</td>
<td>
name to give to the op
</td>
</tr>
</table>



<!-- Tabular view -->
 <table class="responsive fixed orange">
<colgroup><col width="214px"><col></colgroup>
<tr><th colspan="2">Returns</th></tr>

<tr>
<td>
`batch_shape`
</td>
<td>
`Tensor`.
</td>
</tr>
</table>



<h3 id="cdf"><code>cdf</code></h3>

<a target="_blank" href="https://github.com/tensorflow/probability/blob/master/tensorflow_probability/substrates/jax/distributions/distribution.py">View source</a>

<pre class="devsite-click-to-copy prettyprint lang-py tfo-signature-link">
<code>cdf(
    value, name='cdf', **kwargs
)
</code></pre>

Cumulative distribution function.

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

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

Creates a deep copy of the distribution.

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.

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

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.

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

Shannon entropy in nats.

event_shape_tensor

View source

Shape of a single sample from a single batch as a 1-D int32 Tensor.

Args
name name to give to the op

Returns
event_shape Tensor.

is_scalar_batch

View source

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

Indicates that event_shape == [].

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

Returns
is_scalar_event bool scalar Tensor.

kl_divergence

View source

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.

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 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 probability density/mass function.

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
log_prob a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

log_survival_function

View source

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.

mode

View source

Mode.

param_shapes

View source

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

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

Probability density/mass function.

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
prob a Tensor of shape sample_shape(x) + self.batch_shape with values of type self.dtype.

quantile

View source

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.

sample

View source

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.

stddev

View source

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.

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.

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

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