transformed_distribution.py

# Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""A Transformed Distribution class."""

import functools

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.bijectors import ldj_ratio
from tensorflow_probability.python.distributions import batch_broadcast
from tensorflow_probability.python.distributions import distribution as distribution_lib
from tensorflow_probability.python.distributions import kullback_leibler
from tensorflow_probability.python.distributions import log_prob_ratio
from tensorflow_probability.python.internal import auto_composite_tensor
from tensorflow_probability.python.internal import distribution_util as dist_util
from tensorflow_probability.python.internal import parameter_properties
from tensorflow_probability.python.internal import prefer_static as ps
from tensorflow_probability.python.internal import tensorshape_util

__all__ = [
    'TransformedDistribution',
]

JAX_MODE = False


def _default_kwargs_split_fn(kwargs):
  """Default `kwargs` `dict` getter."""
  return (kwargs.get('distribution_kwargs', {}),
          kwargs.get('bijector_kwargs', {}))


class _TransformedDistribution(distribution_lib.Distribution):
  """A Transformed 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`.

  Kullback-Leibler divergence is also well defined for `TransformedDistribution`
  instances that have matching bijectors.  Bijector matching is performed via
  the `Bijector.__eq__` method, e.g., `td1.bijector == td2.bijector`, as part
  of the KL calculation.  If the underlying bijectors do not match, a
  `NotImplementedError` is raised when calling `kl_divergence`.  This is the
  same behavior as calling `kl_divergence` when two distributions do not have
  a registered KL divergence.

  **Note** Due to the current constraints imposed on bijector equality testing,
  `kl_divergence` may behave differently in eager mode computation vs. traced
  computation.  For example, if a TD Bijector's parameters are `Tensor` objects,
  and are themselves derived from e.g. a Variable, some stateful operation, or
  from an argument to a `tf.function` then Bijector equality cannot be known
  during the call to `kl_divergence` and the bijectors are assumed unequal.
  In this case, calling `kl_divergence` may raise an exception in
  graph / tf.function mode, but work just fine in eager / numpy mode.

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

  ```python
  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:

  ```python
  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:

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

  A `TransformedDistribution`'s `batch_shape` is derived by *broadcasting* the
  batch shapes of the base distribution and the bijector. The base distribution
  is then itself implicitly lifted to the broadcast batch shape. For example, in

  ```python
  tfd = tfp.distributions
  tfb = tfp.bijectors
  batch_normal = tfd.TransformedDistribution(
    distribution=tfd.Normal(loc=0., scale=1.),
    bijector=tfb.Shift(shift=[-1., 0., 1.]),
    name='BatchNormalTransformedDistribution')
  ```

  the base distribution has batch shape `[]`, and the bijector applied to this
  distribution contributes a batch shape of `[3]` (obtained as
  `bijector.experimental_batch_shape(
  x_event_ndims=tf.rank(distribution.event_shape))`, yielding the broadcast
  shape `batch_normal.batch_shape == [3]`. Although sampling from the base
  distribution would ordinarily return just a single value, calling
  `batch_normal.sample()` will return a Tensor of 3 independent values, just as
  if the base distribution had explicitly followed the broadcast batch shape.

  The `event_shape` of a `TransformedDistribution` 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`.

  ```python
  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)
  ```

  """

  def __init__(self,
               distribution,
               bijector,
               kwargs_split_fn=_default_kwargs_split_fn,
               validate_args=False,
               parameters=None,
               name=None):
    """Construct a Transformed Distribution.

    Args:
      distribution: The base distribution instance to transform. Typically an
        instance of `Distribution`.
      bijector: The object responsible for calculating the transformation.
        Typically an instance of `Bijector`.
      kwargs_split_fn: 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', {}))`)
      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.
      parameters: Locals dict captured by subclass constructor, to be used for
        copy/slice re-instantiation operations.
      name: Python `str` name prefixed to Ops created by this class. Default:
        `bijector.name + distribution.name`.
    """
    parameters = dict(locals()) if parameters is None else parameters
    name = name or (('' if bijector is None else bijector.name) +
                    (distribution.name or ''))
    with tf.name_scope(name) as name:
      self._distribution = distribution
      self._bijector = bijector
      self._kwargs_split_fn = (_default_kwargs_split_fn
                               if kwargs_split_fn is None
                               else kwargs_split_fn)

      # For convenience we define some handy constants.
      self._zero = tf.constant(0, dtype=tf.int32, name='zero')

      # We don't just want to check isinstance(JointDistribution) because
      # TransformedDistributions with multipart bijectors are effectively
      # joint but don't inherit from JD. The 'duck-type' test is that
      # JDs have a structured dtype.
      dtype = self.bijector.forward_dtype(self.distribution.dtype)
      self._is_joint = tf.nest.is_nested(dtype)

      super(_TransformedDistribution, self).__init__(
          dtype=dtype,
          reparameterization_type=self._distribution.reparameterization_type,
          validate_args=validate_args,
          allow_nan_stats=self._distribution.allow_nan_stats,
          parameters=parameters,
          name=name)

  @property
  def distribution(self):
    """Base distribution, p(x)."""
    return self._distribution

  @property
  def bijector(self):
    """Function transforming x => y."""
    return self._bijector

  @property
  def experimental_is_sharded(self):
    raise NotImplementedError  # TODO(b/175084455): Handle bijector sharding.

  @classmethod
  def _parameter_properties(cls, dtype, num_classes=None):
    return dict(
        distribution=parameter_properties.BatchedComponentProperties(),
        bijector=parameter_properties.BatchedComponentProperties(
            event_ndims=lambda td: tf.nest.map_structure(  # pylint: disable=g-long-lambda
                tensorshape_util.rank, td.distribution.event_shape),
            event_ndims_tensor=lambda td: tf.nest.map_structure(  # pylint: disable=g-long-lambda
                ps.rank_from_shape, td.distribution.event_shape_tensor())))

  def _event_shape_tensor(self):
    return self.bijector.forward_event_shape_tensor(
        self.distribution.event_shape_tensor())

  def _event_shape(self):
    # Since the `bijector` may change the `event_shape`, we then forward what we
    # know to the bijector. This allows the `bijector` to have final say in the
    # `event_shape`.
    return self.bijector.forward_event_shape(self.distribution.event_shape)

  def _batch_shape_tensor(self):
    base_batch_shape_tensor = self.distribution.batch_shape_tensor()
    if tf.nest.is_nested(base_batch_shape_tensor) and self._is_joint:
      # Pass-through rudimentary support for JDs with structured batch shape.
      # TODO(b/194742372): remove support for structured batch shape.
      return tf.nest.pack_sequence_as(
          self.dtype, tf.nest.flatten(base_batch_shape_tensor))
    return super()._batch_shape_tensor()

  def _batch_shape(self):
    batch_shape = self.distribution.batch_shape
    if tf.nest.is_nested(batch_shape) and self._is_joint:
      # Pass-through rudimentary support for JDs with structured batch shape.
      # TODO(b/194742372): remove support for structured batch shape.
      return tf.nest.pack_sequence_as(
          self.dtype, tf.nest.flatten(batch_shape))
    return super()._batch_shape()

  def _maybe_broadcast_distribution_batch_shape(self):
    """Returns the base distribution broadcast to the TD's full batch shape."""
    distribution_batch_shape = self.distribution.batch_shape
    if (tf.nest.is_nested(distribution_batch_shape) or
        tf.nest.is_nested(self.distribution.dtype)):
      # TODO(b/191674464): Support joint distributions in BatchBroadcast.
      return self.distribution

    overall_batch_shape = self.batch_shape
    if (tensorshape_util.is_fully_defined(overall_batch_shape) and
        distribution_batch_shape == overall_batch_shape):
      # No need to broadcast if the distribution already has full batch shape.
      return self.distribution

    if not tensorshape_util.is_fully_defined(overall_batch_shape):
      overall_batch_shape = self.batch_shape_tensor()
    return batch_broadcast.BatchBroadcast(
        self.distribution, with_shape=overall_batch_shape)

  def _call_sample_n(self, sample_shape, seed, **kwargs):
    # We override `_call_sample_n` rather than `_sample_n` so we can ensure that
    # the result of `self.bijector.forward` is not modified (and thus caching
    # works).
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)

    # First, generate samples from the base distribution.
    x = self._maybe_broadcast_distribution_batch_shape().sample(
        sample_shape=sample_shape, seed=seed, **distribution_kwargs)
    # Apply the bijector's forward transformation. For caching to
    # work, it is imperative that this is the last modification to the
    # returned result.
    return self.bijector.forward(x, **bijector_kwargs)

  def _sample_and_log_prob(self, sample_shape, seed, **kwargs):
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      # Computing log_prob with a non-injective bijector requires an explicit
      # inverse to get all points in the inverse image, so we can't get by
      # with just doing the forward pass.
      return super()._sample_and_log_prob(sample_shape, seed=seed, **kwargs)

    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
    x, base_distribution_log_prob = (
        self._maybe_broadcast_distribution_batch_shape(
            ).experimental_sample_and_log_prob(
                sample_shape, seed, **distribution_kwargs))
    y = self.bijector.forward(x, **bijector_kwargs)
    fldj = self.bijector.forward_log_det_jacobian(
        x,
        event_ndims=tf.nest.map_structure(
            ps.rank_from_shape,
            self.distribution.event_shape_tensor()),
        **bijector_kwargs)
    return y, (base_distribution_log_prob -
               tf.cast(fldj, base_distribution_log_prob.dtype))

  def _log_prob(self, y, **kwargs):
    if self.bijector._is_injective:  # pylint: disable=protected-access
      log_prob, _ = self.experimental_local_measure(
          y, backward_compat=True, **kwargs)
      return log_prob

    # TODO(b/197680518): Support base measure handling for non-injective
    # bijectors.
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)

    # For caching to work, it is imperative that the bijector is the first to
    # modify the input.
    x = self.bijector.inverse(y, **bijector_kwargs)
    event_ndims = tf.nest.map_structure(
        ps.rank_from_shape,
        self._event_shape_tensor(),
        self.event_shape)

    ildj = self.bijector.inverse_log_det_jacobian(
        y, event_ndims=event_ndims, **bijector_kwargs)

    # Compute log_prob on each element of the inverse image.
    lp_on_fibers = []
    for x_i, ildj_i in zip(x, ildj):
      base_log_prob = self.distribution.log_prob(x_i, **distribution_kwargs)
      lp_on_fibers.append(base_log_prob + tf.cast(ildj_i, base_log_prob.dtype))
    return tf.reduce_logsumexp(tf.stack(lp_on_fibers), axis=0)

  def _prob(self, y, **kwargs):
    if not hasattr(self.distribution, '_prob') or self.bijector._is_injective:  # pylint: disable=protected-access
      return tf.exp(self._log_prob(y, **kwargs))
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)

    x = self.bijector.inverse(y, **bijector_kwargs)
    event_ndims = tf.nest.map_structure(
        ps.rank_from_shape,
        self._event_shape_tensor(),
        self.event_shape
        )
    ildj = self.bijector.inverse_log_det_jacobian(
        y, event_ndims=event_ndims, **bijector_kwargs)

    # Compute prob on each element of the inverse image.
    prob_on_fibers = []
    for x_i, ildj_i in zip(x, ildj):
      base_prob = self.distribution.prob(x_i, **distribution_kwargs)
      prob_on_fibers.append(
          base_prob * tf.exp(tf.cast(ildj_i, base_prob.dtype)))
    return sum(prob_on_fibers)

  def _log_cdf(self, y, **kwargs):
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError('`log_cdf` is not implemented when '
                                '`bijector` is not injective.')
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
    x = self.bijector.inverse(y, **bijector_kwargs)
    dist = self.distribution
    # TODO(b/141130733): Check/fix any gradient numerics issues.
    return ps.smart_where(
        self.bijector._internal_is_increasing(**bijector_kwargs),  # pylint: disable=protected-access
        lambda: dist.log_cdf(x, **distribution_kwargs),
        lambda: dist.log_survival_function(x, **distribution_kwargs))

  def _cdf(self, y, **kwargs):
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError('`cdf` is not implemented when '
                                '`bijector` is not injective.')
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
    x = self.bijector.inverse(y, **bijector_kwargs)
    # TODO(b/141130733): Check/fix any gradient numerics issues.
    return ps.smart_where(
        self.bijector._internal_is_increasing(**bijector_kwargs),  # pylint: disable=protected-access
        lambda: self.distribution.cdf(x, **distribution_kwargs),
        lambda: self.distribution.survival_function(x, **distribution_kwargs))

  def _log_survival_function(self, y, **kwargs):
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError('`log_survival_function` is not implemented '
                                'when `bijector` is not injective.')
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
    x = self.bijector.inverse(y, **bijector_kwargs)
    dist = self.distribution
    # TODO(b/141130733): Check/fix any gradient numerics issues.
    return ps.smart_where(
        self.bijector._internal_is_increasing(**bijector_kwargs),  # pylint: disable=protected-access
        lambda: dist.log_survival_function(x, **distribution_kwargs),
        lambda: dist.log_cdf(x, **distribution_kwargs))

  def _survival_function(self, y, **kwargs):
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError('`survival_function` is not implemented when '
                                '`bijector` is not injective.')
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
    x = self.bijector.inverse(y, **bijector_kwargs)
    # TODO(b/141130733): Check/fix any gradient numerics issues.
    return ps.smart_where(
        self.bijector._internal_is_increasing(**bijector_kwargs),  # pylint: disable=protected-access
        lambda: self.distribution.survival_function(x, **distribution_kwargs),
        lambda: self.distribution.cdf(x, **distribution_kwargs))

  def _quantile(self, value, **kwargs):
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError('`quantile` is not implemented when '
                                '`bijector` is not injective.')
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
    value = ps.smart_where(
        self.bijector._internal_is_increasing(**bijector_kwargs),  # pylint: disable=protected-access
        lambda: value,
        lambda: 1 - value)
    # x_q is the "qth quantile" of X iff q = P[X <= x_q].  Now, since X =
    # g^{-1}(Y), q = P[X <= x_q] = P[g^{-1}(Y) <= x_q] = P[Y <= g(x_q)],
    # implies the qth quantile of Y is g(x_q).
    inv_cdf = self.distribution.quantile(value, **distribution_kwargs)
    return self.bijector.forward(inv_cdf, **bijector_kwargs)

  def _mode(self, **kwargs):
    return self._mean_mode_impl('mode', kwargs)

  def _mean(self, **kwargs):
    return self._mean_mode_impl('mean', kwargs)

  def _mean_mode_impl(self, attr, kwargs):
    if not self.bijector.is_constant_jacobian:
      raise NotImplementedError(
          f'`{attr}` is not implemented for non-affine `bijectors`.')

    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
    x = getattr(self.distribution, attr)(**distribution_kwargs)
    y = self.bijector.forward(x, **bijector_kwargs)

    sample_shape = tf.convert_to_tensor([], dtype=tf.int32, name='sample_shape')
    y = self._set_sample_static_shape(y, sample_shape, **kwargs)
    return y

  def _stddev(self, **kwargs):
    if not self.bijector.is_constant_jacobian:
      raise NotImplementedError('`stddev` is not implemented for non-affine '
                                '`bijectors`.')
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError('`stddev` is not implemented when '
                                '`bijector` is not injective.')
    if not (self.bijector._is_scalar  # pylint: disable=protected-access
            or self.bijector._is_permutation):  # pylint: disable=protected-access
      raise NotImplementedError('`stddev` is not implemented when `bijector` '
                                'is a multivariate transformation.')

    # A scalar affine bijector is of the form `forward(x) = scale * x + shift`,
    # where the standard deviation is invariant to the shift, so we extract the
    # shift and subtract it.
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
    x_stddev = self.distribution.stddev(**distribution_kwargs)
    y_stddev_plus_shift = self.bijector.forward(x_stddev, **bijector_kwargs)
    shift = self.bijector.forward(
        tf.nest.map_structure(
            tf.zeros_like, x_stddev),
        **bijector_kwargs)
    return tf.nest.map_structure(
        tf.abs,
        tf.nest.map_structure(tf.subtract, y_stddev_plus_shift, shift))

  def _covariance(self, **kwargs):
    if not self.bijector.is_constant_jacobian:
      raise NotImplementedError(
          '`covariance` is not implemented for non-affine `bijectors`.')
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError(
          '`covariance` is not implemented when `bijector` is not injective.')
    if (tf.nest.is_nested(self.bijector.forward_min_event_ndims) or
        self.bijector.forward_event_ndims(1) != 1):
      raise NotImplementedError(
          '`covariance` is only implemented when `bijector` takes vector '
          'inputs and produces vector outputs.')

    # An affine bijector is of the form `forward(x) = scale @ x + shift`,
    # where the covariance is invariant to the shift, so we extract the
    # shift and subtract it.
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
    cov = self.distribution.covariance(**distribution_kwargs)
    zero_vector = tf.zeros_like(cov[..., 0])
    shift = self.bijector.forward(zero_vector, **bijector_kwargs)

    if shift is zero_vector:  # Short-circuit if bijector is tfb.Identity.
      return cov

    # Broadcast `cov` to full batch rank so we can treat its rows as an
    # additional batch dim. Note that we can't just call `forward(cov)` directly
    # because the user presumably lined up the bijector batch dimensions to work
    # when transforming vectors, not matrices.
    cov = tf.broadcast_to(
        cov, ps.broadcast_shape(ps.shape(cov),
                                ps.concat([ps.ones_like(ps.shape(shift)),
                                           [1]], axis=0)))
    ndims = ps.rank(cov)
    cov_rows = dist_util.move_dimension(  # No-op if cov has no batch dims.
        cov, source_idx=-2, dest_idx=0)
    tmp = self.bijector.forward(  # scale @ transpose(cov).
        cov_rows, **bijector_kwargs) - shift
    # Swap leftmost batch dim (rows) with event dim (columns).
    tmp_transpose = tf.transpose(  # cov @ transpose(scale).
        tmp, perm=ps.concat([[ndims - 1], ps.range(1, ndims - 1), [0]], axis=0))
    result_rows = self.bijector.forward(  # scale @ cov @ transpose(scale).
        tmp_transpose, **bijector_kwargs) - shift
    return dist_util.move_dimension(  # No-op if result has no batch dims.
        result_rows, source_idx=0, dest_idx=-2)

  def _entropy(self, **kwargs):
    if not self.bijector.is_constant_jacobian:
      raise NotImplementedError('`entropy` is not implemented.')
    if not self.bijector._is_injective:  # pylint: disable=protected-access
      raise NotImplementedError('`entropy` is not implemented when '
                                '`bijector` is not injective.')
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
    # Suppose Y = g(X) where g is a diffeomorphism and X is a continuous rv. It
    # can be shown that:
    #   H[Y] = H[X] + E_X[(log o abs o det o J o g)(X)].
    # If is_constant_jacobian then:
    #   E_X[(log o abs o det o J o g)(X)] = (log o abs o det o J o g)(c)
    # where c can by anything.
    entropy = self.distribution.entropy(**distribution_kwargs)

    # Create a dummy event of zeros to pass to
    # `bijector.inverse_log_det_jacobian` to extract the constant Jacobian.
    event_shape_tensor = self._event_shape_tensor()
    event_ndims = tf.nest.map_structure(
        ps.rank_from_shape,
        event_shape_tensor, self.event_shape)
    dummy = tf.nest.map_structure(
        ps.zeros, event_shape_tensor, self.dtype)

    ildj = self.bijector.inverse_log_det_jacobian(
        dummy, event_ndims=event_ndims, **bijector_kwargs)

    entropy = entropy - tf.cast(ildj, entropy.dtype)
    tensorshape_util.set_shape(entropy, self.batch_shape)
    return entropy

  # pylint: disable=not-callable
  def _default_event_space_bijector(self):
    if self.distribution.experimental_default_event_space_bijector() is None:
      return None
    return self.bijector(
        self.distribution.experimental_default_event_space_bijector())
  # pylint: enable=not-callable

  def experimental_local_measure(self, y, backward_compat=False, **kwargs):
    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)

    # For caching to work, it is imperative that the bijector is the first to
    # modify the input.
    x = self.bijector.inverse(y, **bijector_kwargs)
    event_ndims = self.bijector.inverse_event_ndims(
        tf.nest.map_structure(ps.rank_from_shape, self._event_shape_tensor(),
                              self.event_shape), **bijector_kwargs)

    if self.bijector._is_injective:  # pylint: disable=protected-access
      local_measure_fn = self.distribution.experimental_local_measure
      density_corr_fn = self.bijector.experimental_compute_density_correction
      base_log_prob, tangent_space = local_measure_fn(
          x, backward_compat=backward_compat, **distribution_kwargs)
      correction, new_tangent_space = density_corr_fn(
          x,
          tangent_space,
          backward_compat=backward_compat,
          event_ndims=event_ndims,
          **bijector_kwargs)
      log_prob = base_log_prob - tf.cast(correction, base_log_prob.dtype)
      return log_prob, new_tangent_space
    else:
      raise NotImplementedError


class _TransformedDistributionMeta(distribution_lib._DistributionMeta):  # pylint: disable=protected-access
  """Metaclass for TransformedDistribution.

  This metaclass ensures that subclasses of TransformedDistribution are
  AutoCompositeTensors. TransformedDistribution itself is not an
  AutoCompositeTensor, since we define its type spec differently depending on
  whether `split_kwargs_fn` is the default (in which case we omit it from the
  type spec, making the spec serializable).
  """

  def __new__(mcs, classname, baseclasses, attrs):  # pylint: disable=bad-classmethod-argument
    cls = super(_TransformedDistributionMeta, mcs).__new__(
        mcs, classname, baseclasses, attrs)
    if ((cls.__module__ ==
         'tensorflow_probability.python.distributions.transformed_distribution')
        and classname == 'TransformedDistribution'):
      return cls

    # Ensure that subclasses of TransformedDistribution are
    # AutoCompositeTensors.
    if 'tensorflow_probability.python.distributions' in cls.__module__:
      module_name = 'tfp.distributions'
    elif ('tensorflow_probability.python.experimental.distributions'
          in cls.__module__):
      module_name = 'tfp.experimental.distributions'
    else:
      module_name = cls.__module__
    return auto_composite_tensor.auto_composite_tensor(
        cls,
        omit_kwargs=('parameters',),
        non_identifying_kwargs=('name',),
        module_name=module_name)


class TransformedDistribution(
    _TransformedDistribution, tf.__internal__.CompositeTensor,
    metaclass=_TransformedDistributionMeta):

  def __new__(cls, *args, **kwargs):
    """Maybe return a non-`CompositeTensor` `_TransformedDistribution`."""

    if cls is TransformedDistribution:
      if args:
        distribution = args[0]
      else:
        distribution = kwargs.get('distribution')
      if len(args) > 1:
        bijector = args[1]
      else:
        bijector = kwargs.get('bijector')

      if not (auto_composite_tensor.is_composite_tensor(distribution)
              and auto_composite_tensor.is_composite_tensor(bijector)):
        return _TransformedDistribution(*args, **kwargs)
    return super(TransformedDistribution, cls).__new__(cls)

  @property
  def _type_spec(self):
    # If `kwargs_split_fn` is the default, omit it so the type spec is
    # serializable.
    if self._kwargs_split_fn is _default_kwargs_split_fn:
      omit_kwargs = ('parameters', 'kwargs_split_fn')
    else:
      omit_kwargs = ('parameters',)
    return _TransformedDistributionSpec.from_instance(
        self, omit_kwargs=omit_kwargs, non_identifying_kwargs=('name',))

  def _convert_variables_to_tensors(self):
    return auto_composite_tensor.convert_variables_to_tensors(self)


@auto_composite_tensor.type_spec_register(
    'tfp.distributions.TransformedDistributionSpec')
class _TransformedDistributionSpec(
    auto_composite_tensor._AutoCompositeTensorTypeSpec):  # pylint: disable=protected-access

  @property
  def value_type(self):
    return TransformedDistribution


if JAX_MODE:
  from jax import tree_util  # pylint: disable=g-import-not-at-top
  tree_util.register_pytree_node(
      TransformedDistribution,
      auto_composite_tensor.pytree_flatten,
      functools.partial(auto_composite_tensor.pytree_unflatten,
                        TransformedDistribution))


TransformedDistribution.__doc__ = _TransformedDistribution.__doc__ + '\n' + (
    'If both `distribution` and `bijector` are `CompositeTensor`s, then the '
    'resulting `TransformedDistribution` instance is a `CompositeTensor` as '
    'well. Otherwise, a non-`CompositeTensor` `_TransformedDistribution` '
    'instance is created instead. Distribution subclasses that inherit from '
    '`TransformedDistribution` will also inherit from `CompositeTensor`.')


@kullback_leibler.RegisterKL(
    _TransformedDistribution, _TransformedDistribution)
def _kl_transformed_transformed(a, b, name=None):
  """Calculate the batched KL divergence KL(a || b) with a and b Transformed.

  Args:
    a: instance of a TransformedDistribution object.
    b: instance of a TransformedDistribution object.
    name: Name to use for created operations.
      Default value: `None` (i.e., `'kl_normal_normal'`).

  Returns:
    kl_div: Batchwise KL(a || b)

  Raises:
    NotImplementedError: If `a.bijector != b.bijector`.
  """
  with tf.name_scope(name or 'kl_transformed_transformed'):
    if a.bijector == b.bijector:
      return kullback_leibler.kl_divergence(a.distribution, b.distribution)
  raise NotImplementedError(
      'Unable to calculate KL divergence between {} and {} because '
      'their bijectors are not equal: {} vs. {}'.format(
          a, b, a.bijector, b.bijector))


@log_prob_ratio.RegisterLogProbRatio(_TransformedDistribution)
def _transformed_log_prob_ratio(p, x, q, y, name=None):
  """Computes p.log_prob(x) - q.log_prob(y) for p and q both TDs."""
  with tf.name_scope(name or 'transformed_log_prob_ratio'):
    x_ = p.bijector.inverse(x)
    y_ = q.bijector.inverse(y)

    base_log_prob_ratio = log_prob_ratio.log_prob_ratio(
        p.distribution, x_, q.distribution, y_)

    event_ndims = tf.nest.map_structure(
        ps.rank_from_shape,
        p.event_shape_tensor,
        tf.nest.map_structure(tensorshape_util.merge_with,
                              p.event_shape, q.event_shape))
    ildj_ratio = ldj_ratio.inverse_log_det_jacobian_ratio(
        p.bijector, x, q.bijector, y, event_ndims)
    return base_log_prob_ratio + tf.cast(ildj_ratio, base_log_prob_ratio.dtype)