Table of Contents
Shortcuts

# Source code for torch.distributions.transforms

    import functools
    import math
    import numbers
    import operator
    import weakref
    from typing import List

    import torch
    import torch.nn.functional as F
    from torch.distributions import constraints
    from torch.distributions.utils import (
        _sum_rightmost,
        broadcast_all,
        lazy_property,
        tril_matrix_to_vec,
        vec_to_tril_matrix,
    )
    from torch.nn.functional import pad, softplus

    __all__ = [
        "AbsTransform",
        "AffineTransform",
        "CatTransform",
        "ComposeTransform",
        "CorrCholeskyTransform",
        "CumulativeDistributionTransform",
        "ExpTransform",
        "IndependentTransform",
        "LowerCholeskyTransform",
        "PositiveDefiniteTransform",
        "PowerTransform",
        "ReshapeTransform",
        "SigmoidTransform",
        "SoftplusTransform",
        "TanhTransform",
        "SoftmaxTransform",
        "StackTransform",
        "StickBreakingTransform",
        "Transform",
        "identity_transform",
    ]

    class Transform:
        """
        Abstract class for invertable transformations with computable log
        det jacobians. They are primarily used in
        :class:`torch.distributions.TransformedDistribution`.

        Caching is useful for transforms whose inverses are either expensive or
        numerically unstable. Note that care must be taken with memoized values
        since the autograd graph may be reversed. For example while the following
        works with or without caching::

            y = t(x)
            t.log_abs_det_jacobian(x, y).backward()  # x will receive gradients.

        However the following will error when caching due to dependency reversal::

            y = t(x)
            z = t.inv(y)
            grad(z.sum(), [y])  # error because z is x

        Derived classes should implement one or both of :meth:`_call` or
        :meth:`_inverse`. Derived classes that set `bijective=True` should also
        implement :meth:`log_abs_det_jacobian`.

        Args:
            cache_size (int): Size of cache. If zero, no caching is done. If one,
                the latest single value is cached. Only 0 and 1 are supported.

        Attributes:
            domain (:class:`~torch.distributions.constraints.Constraint`):
                The constraint representing valid inputs to this transform.
            codomain (:class:`~torch.distributions.constraints.Constraint`):
                The constraint representing valid outputs to this transform
                which are inputs to the inverse transform.
            bijective (bool): Whether this transform is bijective. A transform
                ``t`` is bijective iff ``t.inv(t(x)) == x`` and
                ``t(t.inv(y)) == y`` for every ``x`` in the domain and ``y`` in
                the codomain. Transforms that are not bijective should at least
                maintain the weaker pseudoinverse properties
                ``t(t.inv(t(x)) == t(x)`` and ``t.inv(t(t.inv(y))) == t.inv(y)``.
            sign (int or Tensor): For bijective univariate transforms, this
                should be +1 or -1 depending on whether transform is monotone
                increasing or decreasing.
        """

        bijective = False
        domain: constraints.Constraint
        codomain: constraints.Constraint

        def __init__(self, cache_size=0):
            self._cache_size = cache_size
            self._inv = None
            if cache_size == 0:
                pass  # default behavior
            elif cache_size == 1:
                self._cached_x_y = None, None
            else:
                raise ValueError("cache_size must be 0 or 1")
            super().__init__()

        def __getstate__(self):
            state = self.__dict__.copy()
            state["_inv"] = None
            return state

        @property
        def event_dim(self):
            if self.domain.event_dim == self.codomain.event_dim:
                return self.domain.event_dim
            raise ValueError("Please use either .domain.event_dim or .codomain.event_dim")

        @property
        def inv(self):
            """
            Returns the inverse :class:`Transform` of this transform.
            This should satisfy ``t.inv.inv is t``.
            """
            inv = None
            if self._inv is not None:
                inv = self._inv()
            if inv is None:
                inv = _InverseTransform(self)
                self._inv = weakref.ref(inv)
            return inv

        @property
        def sign(self):
            """
            Returns the sign of the determinant of the Jacobian, if applicable.
            In general this only makes sense for bijective transforms.
            """
            raise NotImplementedError

        def with_cache(self, cache_size=1):
            if self._cache_size == cache_size:
                return self
            if type(self).__init__ is Transform.__init__:
                return type(self)(cache_size=cache_size)
            raise NotImplementedError(f"{type(self)}.with_cache is not implemented")

        def __eq__(self, other):
            return self is other

        def __ne__(self, other):
            # Necessary for Python2
            return not self.__eq__(other)

        def __call__(self, x):
            """
            Computes the transform `x => y`.
            """
            if self._cache_size == 0:
                return self._call(x)
            x_old, y_old = self._cached_x_y
            if x is x_old:
                return y_old
            y = self._call(x)
            self._cached_x_y = x, y
            return y

        def _inv_call(self, y):
            """
            Inverts the transform `y => x`.
            """
            if self._cache_size == 0:
                return self._inverse(y)
            x_old, y_old = self._cached_x_y
            if y is y_old:
                return x_old
            x = self._inverse(y)
            self._cached_x_y = x, y
            return x

        def _call(self, x):
            """
            Abstract method to compute forward transformation.
            """
            raise NotImplementedError

        def _inverse(self, y):
            """
            Abstract method to compute inverse transformation.
            """
            raise NotImplementedError

        def log_abs_det_jacobian(self, x, y):
            """
            Computes the log det jacobian `log |dy/dx|` given input and output.
            """
            raise NotImplementedError

        def __repr__(self):
            return self.__class__.__name__ + "()"

        def forward_shape(self, shape):
            """
            Infers the shape of the forward computation, given the input shape.
            Defaults to preserving shape.
            """
            return shape

        def inverse_shape(self, shape):
            """
            Infers the shapes of the inverse computation, given the output shape.
            Defaults to preserving shape.
            """
            return shape

    class _InverseTransform(Transform):
        """
        Inverts a single :class:`Transform`.
        This class is private; please instead use the ``Transform.inv`` property.
        """

        def __init__(self, transform: Transform):
            super().__init__(cache_size=transform._cache_size)
            self._inv: Transform = transform

        @constraints.dependent_property(is_discrete=False)
        def domain(self):
            assert self._inv is not None
            return self._inv.codomain

        @constraints.dependent_property(is_discrete=False)
        def codomain(self):
            assert self._inv is not None
            return self._inv.domain

        @property
        def bijective(self):
            assert self._inv is not None
            return self._inv.bijective

        @property
        def sign(self):
            assert self._inv is not None
            return self._inv.sign

        @property
        def inv(self):
            return self._inv

        def with_cache(self, cache_size=1):
            assert self._inv is not None
            return self.inv.with_cache(cache_size).inv

        def __eq__(self, other):
            if not isinstance(other, _InverseTransform):
                return False
            assert self._inv is not None
            return self._inv == other._inv

        def __repr__(self):
            return f"{self.__class__.__name__}({repr(self._inv)})"

        def __call__(self, x):
            assert self._inv is not None
            return self._inv._inv_call(x)

        def log_abs_det_jacobian(self, x, y):
            assert self._inv is not None
            return -self._inv.log_abs_det_jacobian(y, x)

        def forward_shape(self, shape):
            return self._inv.inverse_shape(shape)

        def inverse_shape(self, shape):
            return self._inv.forward_shape(shape)

    class ComposeTransform(Transform):
        """
        Composes multiple transforms in a chain.
        The transforms being composed are responsible for caching.

        Args:
            parts (list of :class:`Transform`): A list of transforms to compose.
            cache_size (int): Size of cache. If zero, no caching is done. If one,
                the latest single value is cached. Only 0 and 1 are supported.
        """

        def __init__(self, parts: List[Transform], cache_size=0):
            if cache_size:
                parts = [part.with_cache(cache_size) for part in parts]
            super().__init__(cache_size=cache_size)
            self.parts = parts

        def __eq__(self, other):
            if not isinstance(other, ComposeTransform):
                return False
            return self.parts == other.parts

        @constraints.dependent_property(is_discrete=False)
        def domain(self):
            if not self.parts:
                return constraints.real
            domain = self.parts[0].domain
            # Adjust event_dim to be maximum among all parts.
            event_dim = self.parts[-1].codomain.event_dim
            for part in reversed(self.parts):
                event_dim += part.domain.event_dim - part.codomain.event_dim
                event_dim = max(event_dim, part.domain.event_dim)
            assert event_dim >= domain.event_dim
            if event_dim > domain.event_dim:
                domain = constraints.independent(domain, event_dim - domain.event_dim)
            return domain

        @constraints.dependent_property(is_discrete=False)
        def codomain(self):
            if not self.parts:
                return constraints.real
            codomain = self.parts[-1].codomain
            # Adjust event_dim to be maximum among all parts.
            event_dim = self.parts[0].domain.event_dim
            for part in self.parts:
                event_dim += part.codomain.event_dim - part.domain.event_dim
                event_dim = max(event_dim, part.codomain.event_dim)
            assert event_dim >= codomain.event_dim
            if event_dim > codomain.event_dim:
                codomain = constraints.independent(codomain, event_dim - codomain.event_dim)
            return codomain

        @lazy_property
        def bijective(self):
            return all(p.bijective for p in self.parts)

        @lazy_property
        def sign(self):
            sign = 1
            for p in self.parts:
                sign = sign * p.sign
            return sign

        @property
        def inv(self):
            inv = None
            if self._inv is not None:
                inv = self._inv()
            if inv is None:
                inv = ComposeTransform([p.inv for p in reversed(self.parts)])
                self._inv = weakref.ref(inv)
                inv._inv = weakref.ref(self)
            return inv

        def with_cache(self, cache_size=1):
            if self._cache_size == cache_size:
                return self
            return ComposeTransform(self.parts, cache_size=cache_size)

        def __call__(self, x):
            for part in self.parts:
                x = part(x)
            return x

        def log_abs_det_jacobian(self, x, y):
            if not self.parts:
                return torch.zeros_like(x)

            # Compute intermediates. This will be free if parts[:-1] are all cached.
            xs = [x]
            for part in self.parts[:-1]:
                xs.append(part(xs[-1]))
            xs.append(y)

            terms = []
            event_dim = self.domain.event_dim
            for part, x, y in zip(self.parts, xs[:-1], xs[1:]):
                terms.append(
                    _sum_rightmost(
                        part.log_abs_det_jacobian(x, y), event_dim - part.domain.event_dim
                    )
                )
                event_dim += part.codomain.event_dim - part.domain.event_dim
            return functools.reduce(operator.add, terms)

        def forward_shape(self, shape):
            for part in self.parts:
                shape = part.forward_shape(shape)
            return shape

        def inverse_shape(self, shape):
            for part in reversed(self.parts):
                shape = part.inverse_shape(shape)
            return shape

        def __repr__(self):
            fmt_string = self.__class__.__name__ + "(\n    "
            fmt_string += ",\n    ".join([p.__repr__() for p in self.parts])
            fmt_string += "\n)"
            return fmt_string

    identity_transform = ComposeTransform([])

    class IndependentTransform(Transform):
        """
        Wrapper around another transform to treat
        ``reinterpreted_batch_ndims``-many extra of the right most dimensions as
        dependent. This has no effect on the forward or backward transforms, but
        does sum out ``reinterpreted_batch_ndims``-many of the rightmost dimensions
        in :meth:`log_abs_det_jacobian`.

        Args:
            base_transform (:class:`Transform`): A base transform.
            reinterpreted_batch_ndims (int): The number of extra rightmost
                dimensions to treat as dependent.
        """

        def __init__(self, base_transform, reinterpreted_batch_ndims, cache_size=0):
            super().__init__(cache_size=cache_size)
            self.base_transform = base_transform.with_cache(cache_size)
            self.reinterpreted_batch_ndims = reinterpreted_batch_ndims

        def with_cache(self, cache_size=1):
            if self._cache_size == cache_size:
                return self
            return IndependentTransform(
                self.base_transform, self.reinterpreted_batch_ndims, cache_size=cache_size
            )

        @constraints.dependent_property(is_discrete=False)
        def domain(self):
            return constraints.independent(
                self.base_transform.domain, self.reinterpreted_batch_ndims
            )

        @constraints.dependent_property(is_discrete=False)
        def codomain(self):
            return constraints.independent(
                self.base_transform.codomain, self.reinterpreted_batch_ndims
            )

        @property
        def bijective(self):
            return self.base_transform.bijective

        @property
        def sign(self):
            return self.base_transform.sign

        def _call(self, x):
            if x.dim() < self.domain.event_dim:
                raise ValueError("Too few dimensions on input")
            return self.base_transform(x)

        def _inverse(self, y):
            if y.dim() < self.codomain.event_dim:
                raise ValueError("Too few dimensions on input")
            return self.base_transform.inv(y)

        def log_abs_det_jacobian(self, x, y):
            result = self.base_transform.log_abs_det_jacobian(x, y)
            result = _sum_rightmost(result, self.reinterpreted_batch_ndims)
            return result

        def __repr__(self):
            return f"{self.__class__.__name__}({repr(self.base_transform)}, {self.reinterpreted_batch_ndims})"

        def forward_shape(self, shape):
            return self.base_transform.forward_shape(shape)

        def inverse_shape(self, shape):
            return self.base_transform.inverse_shape(shape)

    class ReshapeTransform(Transform):
        """
        Unit Jacobian transform to reshape the rightmost part of a tensor.

        Note that ``in_shape`` and ``out_shape`` must have the same number of
        elements, just as for :meth:`torch.Tensor.reshape`.

        Arguments:
            in_shape (torch.Size): The input event shape.
            out_shape (torch.Size): The output event shape.
        """

        bijective = True

        def __init__(self, in_shape, out_shape, cache_size=0):
            self.in_shape = torch.Size(in_shape)
            self.out_shape = torch.Size(out_shape)
            if self.in_shape.numel() != self.out_shape.numel():
                raise ValueError("in_shape, out_shape have different numbers of elements")
            super().__init__(cache_size=cache_size)

        @constraints.dependent_property
        def domain(self):
            return constraints.independent(constraints.real, len(self.in_shape))

        @constraints.dependent_property
        def codomain(self):
            return constraints.independent(constraints.real, len(self.out_shape))

        def with_cache(self, cache_size=1):
            if self._cache_size == cache_size:
                return self
            return ReshapeTransform(self.in_shape, self.out_shape, cache_size=cache_size)

        def _call(self, x):
            batch_shape = x.shape[: x.dim() - len(self.in_shape)]
            return x.reshape(batch_shape + self.out_shape)

        def _inverse(self, y):
            batch_shape = y.shape[: y.dim() - len(self.out_shape)]
            return y.reshape(batch_shape + self.in_shape)

        def log_abs_det_jacobian(self, x, y):
            batch_shape = x.shape[: x.dim() - len(self.in_shape)]
            return x.new_zeros(batch_shape)

        def forward_shape(self, shape):
            if len(shape) < len(self.in_shape):
                raise ValueError("Too few dimensions on input")
            cut = len(shape) - len(self.in_shape)
            if shape[cut:] != self.in_shape:
                raise ValueError(
                    f"Shape mismatch: expected {shape[cut:]} but got {self.in_shape}"
                )
            return shape[:cut] + self.out_shape

        def inverse_shape(self, shape):
            if len(shape) < len(self.out_shape):
                raise ValueError("Too few dimensions on input")
            cut = len(shape) - len(self.out_shape)
            if shape[cut:] != self.out_shape:
                raise ValueError(
                    f"Shape mismatch: expected {shape[cut:]} but got {self.out_shape}"
                )
            return shape[:cut] + self.in_shape

    class ExpTransform(Transform):
        r"""
        Transform via the mapping :math:`y = \exp(x)`.
        """
        domain = constraints.real
        codomain = constraints.positive
        bijective = True
        sign = +1

        def __eq__(self, other):
            return isinstance(other, ExpTransform)

        def _call(self, x):
            return x.exp()

        def _inverse(self, y):
            return y.log()

        def log_abs_det_jacobian(self, x, y):
            return x

    class PowerTransform(Transform):
        r"""
        Transform via the mapping :math:`y = x^{\text{exponent}}`.
        """
        domain = constraints.positive
        codomain = constraints.positive
        bijective = True
        sign = +1

        def __init__(self, exponent, cache_size=0):
            super().__init__(cache_size=cache_size)
            (self.exponent,) = broadcast_all(exponent)

        def with_cache(self, cache_size=1):
            if self._cache_size == cache_size:
                return self
            return PowerTransform(self.exponent, cache_size=cache_size)

        def __eq__(self, other):
            if not isinstance(other, PowerTransform):
                return False
            return self.exponent.eq(other.exponent).all().item()

        def _call(self, x):
            return x.pow(self.exponent)

        def _inverse(self, y):
            return y.pow(1 / self.exponent)

        def log_abs_det_jacobian(self, x, y):
            return (self.exponent * y / x).abs().log()

        def forward_shape(self, shape):
            return torch.broadcast_shapes(shape, getattr(self.exponent, "shape", ()))

        def inverse_shape(self, shape):
            return torch.broadcast_shapes(shape, getattr(self.exponent, "shape", ()))

    def _clipped_sigmoid(x):
        finfo = torch.finfo(x.dtype)
        return torch.clamp(torch.sigmoid(x), min=finfo.tiny, max=1.0 - finfo.eps)

    class SigmoidTransform(Transform):
        r"""
        Transform via the mapping :math:`y = \frac{1}{1 + \exp(-x)}` and :math:`x = \text{logit}(y)`.
        """
        domain = constraints.real
        codomain = constraints.unit_interval
        bijective = True
        sign = +1

        def __eq__(self, other):
            return isinstance(other, SigmoidTransform)

        def _call(self, x):
            return _clipped_sigmoid(x)

        def _inverse(self, y):
            finfo = torch.finfo(y.dtype)
            y = y.clamp(min=finfo.tiny, max=1.0 - finfo.eps)
            return y.log() - (-y).log1p()

        def log_abs_det_jacobian(self, x, y):
            return -F.softplus(-x) - F.softplus(x)

    class SoftplusTransform(Transform):
        r"""
        Transform via the mapping :math:`\text{Softplus}(x) = \log(1 + \exp(x))`.
        The implementation reverts to the linear function when :math:`x > 20`.
        """
        domain = constraints.real
        codomain = constraints.positive
        bijective = True
        sign = +1

        def __eq__(self, other):
            return isinstance(other, SoftplusTransform)

        def _call(self, x):
            return softplus(x)

        def _inverse(self, y):
            return (-y).expm1().neg().log() + y

        def log_abs_det_jacobian(self, x, y):
            return -softplus(-x)

    class TanhTransform(Transform):
        r"""
        Transform via the mapping :math:`y = \tanh(x)`.

        It is equivalent to
        ```
        ComposeTransform([AffineTransform(0., 2.), SigmoidTransform(), AffineTransform(-1., 2.)])
        ```
        However this might not be numerically stable, thus it is recommended to use `TanhTransform`
        instead.

        Note that one should use `cache_size=1` when it comes to `NaN/Inf` values.

        """
        domain = constraints.real
        codomain = constraints.interval(-1.0, 1.0)
        bijective = True
        sign = +1

        def __eq__(self, other):
            return isinstance(other, TanhTransform)

        def _call(self, x):
            return x.tanh()

        def _inverse(self, y):
            # We do not clamp to the boundary here as it may degrade the performance of certain algorithms.
            # one should use `cache_size=1` instead
            return torch.atanh(y)

        def log_abs_det_jacobian(self, x, y):
            # We use a formula that is more numerically stable, see details in the following link
            # https://github.com/tensorflow/probability/blob/master/tensorflow_probability/python/bijectors/tanh.py#L69-L80
            return 2.0 * (math.log(2.0) - x - softplus(-2.0 * x))

    class AbsTransform(Transform):
        r"""
        Transform via the mapping :math:`y = |x|`.
        """
        domain = constraints.real
        codomain = constraints.positive

        def __eq__(self, other):
            return isinstance(other, AbsTransform)

        def _call(self, x):
            return x.abs()

        def _inverse(self, y):
            return y

    class AffineTransform(Transform):
        r"""
        Transform via the pointwise affine mapping :math:`y = \text{loc} + \text{scale} \times x`.

        Args:
            loc (Tensor or float): Location parameter.
            scale (Tensor or float): Scale parameter.
            event_dim (int): Optional size of `event_shape`. This should be zero
                for univariate random variables, 1 for distributions over vectors,
                2 for distributions over matrices, etc.
        """
        bijective = True

        def __init__(self, loc, scale, event_dim=0, cache_size=0):
            super().__init__(cache_size=cache_size)
            self.loc = loc
            self.scale = scale
            self._event_dim = event_dim

        @property
        def event_dim(self):
            return self._event_dim

        @constraints.dependent_property(is_discrete=False)
        def domain(self):
            if self.event_dim == 0:
                return constraints.real
            return constraints.independent(constraints.real, self.event_dim)

        @constraints.dependent_property(is_discrete=False)
        def codomain(self):
            if self.event_dim == 0:
                return constraints.real
            return constraints.independent(constraints.real, self.event_dim)

        def with_cache(self, cache_size=1):
            if self._cache_size == cache_size:
                return self
            return AffineTransform(
                self.loc, self.scale, self.event_dim, cache_size=cache_size
            )

        def __eq__(self, other):
            if not isinstance(other, AffineTransform):
                return False

            if isinstance(self.loc, numbers.Number) and isinstance(
                other.loc, numbers.Number
            ):
                if self.loc != other.loc:
                    return False
            else:
                if not (self.loc == other.loc).all().item():
                    return False

            if isinstance(self.scale, numbers.Number) and isinstance(
                other.scale, numbers.Number
            ):
                if self.scale != other.scale:
                    return False
            else:
                if not (self.scale == other.scale).all().item():
                    return False

            return True

        @property
        def sign(self):
            if isinstance(self.scale, numbers.Real):
                return 1 if float(self.scale) > 0 else -1 if float(self.scale) < 0 else 0
            return self.scale.sign()

        def _call(self, x):
            return self.loc + self.scale * x

        def _inverse(self, y):
            return (y - self.loc) / self.scale

        def log_abs_det_jacobian(self, x, y):
            shape = x.shape
            scale = self.scale
            if isinstance(scale, numbers.Real):
                result = torch.full_like(x, math.log(abs(scale)))
            else:
                result = torch.abs(scale).log()
            if self.event_dim:
                result_size = result.size()[: -self.event_dim] + (-1,)
                result = result.view(result_size).sum(-1)
                shape = shape[: -self.event_dim]
            return result.expand(shape)

        def forward_shape(self, shape):
            return torch.broadcast_shapes(
                shape, getattr(self.loc, "shape", ()), getattr(self.scale, "shape", ())
            )

        def inverse_shape(self, shape):
            return torch.broadcast_shapes(
                shape, getattr(self.loc, "shape", ()), getattr(self.scale, "shape", ())
            )

    class CorrCholeskyTransform(Transform):
        r"""
        Transforms an uncontrained real vector :math:`x` with length :math:`D*(D-1)/2` into the
        Cholesky factor of a D-dimension correlation matrix. This Cholesky factor is a lower
        triangular matrix with positive diagonals and unit Euclidean norm for each row.
        The transform is processed as follows:

            1. First we convert x into a lower triangular matrix in row order.
            2. For each row :math:`X_i` of the lower triangular part, we apply a *signed* version of
               class :class:`StickBreakingTransform` to transform :math:`X_i` into a
               unit Euclidean length vector using the following steps:
               - Scales into the interval :math:`(-1, 1)` domain: :math:`r_i = \tanh(X_i)`.
               - Transforms into an unsigned domain: :math:`z_i = r_i^2`.
               - Applies :math:`s_i = StickBreakingTransform(z_i)`.
               - Transforms back into signed domain: :math:`y_i = sign(r_i) * \sqrt{s_i}`.
        """
        domain = constraints.real_vector
        codomain = constraints.corr_cholesky
        bijective = True

        def _call(self, x):
            x = torch.tanh(x)
            eps = torch.finfo(x.dtype).eps
            x = x.clamp(min=-1 + eps, max=1 - eps)
            r = vec_to_tril_matrix(x, diag=-1)
            # apply stick-breaking on the squared values
            # Note that y = sign(r) * sqrt(z * z1m_cumprod)
            #             = (sign(r) * sqrt(z)) * sqrt(z1m_cumprod) = r * sqrt(z1m_cumprod)
            z = r**2
            z1m_cumprod_sqrt = (1 - z).sqrt().cumprod(-1)
            # Diagonal elements must be 1.
            r = r + torch.eye(r.shape[-1], dtype=r.dtype, device=r.device)
            y = r * pad(z1m_cumprod_sqrt[..., :-1], [1, 0], value=1)
            return y

        def _inverse(self, y):
            # inverse stick-breaking
            # See: https://mc-stan.org/docs/2_18/reference-manual/cholesky-factors-of-correlation-matrices-1.html
            y_cumsum = 1 - torch.cumsum(y * y, dim=-1)
            y_cumsum_shifted = pad(y_cumsum[..., :-1], [1, 0], value=1)
            y_vec = tril_matrix_to_vec(y, diag=-1)
            y_cumsum_vec = tril_matrix_to_vec(y_cumsum_shifted, diag=-1)
            t = y_vec / (y_cumsum_vec).sqrt()
            # inverse of tanh
            x = (t.log1p() - t.neg().log1p()) / 2
            return x

        def log_abs_det_jacobian(self, x, y, intermediates=None):
            # Because domain and codomain are two spaces with different dimensions, determinant of
            # Jacobian is not well-defined. We return `log_abs_det_jacobian` of `x` and the
            # flattened lower triangular part of `y`.

            # See: https://mc-stan.org/docs/2_18/reference-manual/cholesky-factors-of-correlation-matrices-1.html
            y1m_cumsum = 1 - (y * y).cumsum(dim=-1)
            # by taking diagonal=-2, we don't need to shift z_cumprod to the right
            # also works for 2 x 2 matrix
            y1m_cumsum_tril = tril_matrix_to_vec(y1m_cumsum, diag=-2)
            stick_breaking_logdet = 0.5 * (y1m_cumsum_tril).log().sum(-1)
            tanh_logdet = -2 * (x + softplus(-2 * x) - math.log(2.0)).sum(dim=-1)
            return stick_breaking_logdet + tanh_logdet

        def forward_shape(self, shape):
            # Reshape from (..., N) to (..., D, D).
            if len(shape) < 1:
                raise ValueError("Too few dimensions on input")
            N = shape[-1]
            D = round((0.25 + 2 * N) ** 0.5 + 0.5)
            if D * (D - 1) // 2 != N:
                raise ValueError("Input is not a flattend lower-diagonal number")
            return shape[:-1] + (D, D)

        def inverse_shape(self, shape):
            # Reshape from (..., D, D) to (..., N).
            if len(shape) < 2:
                raise ValueError("Too few dimensions on input")
            if shape[-2] != shape[-1]:
                raise ValueError("Input is not square")
            D = shape[-1]
            N = D * (D - 1) // 2
            return shape[:-2] + (N,)

    class SoftmaxTransform(Transform):
        r"""
        Transform from unconstrained space to the simplex via :math:`y = \exp(x)` then
        normalizing.

        This is not bijective and cannot be used for HMC. However this acts mostly
        coordinate-wise (except for the final normalization), and thus is
        appropriate for coordinate-wise optimization algorithms.
        """
        domain = constraints.real_vector
        codomain = constraints.simplex

        def __eq__(self, other):
            return isinstance(other, SoftmaxTransform)

        def _call(self, x):
            logprobs = x
            probs = (logprobs - logprobs.max(-1, True)[0]).exp()
            return probs / probs.sum(-1, True)

        def _inverse(self, y):
            probs = y
            return probs.log()

        def forward_shape(self, shape):
            if len(shape) < 1:
                raise ValueError("Too few dimensions on input")
            return shape

        def inverse_shape(self, shape):
            if len(shape) < 1:
                raise ValueError("Too few dimensions on input")
            return shape

    class StickBreakingTransform(Transform):
        """
        Transform from unconstrained space to the simplex of one additional
        dimension via a stick-breaking process.

        This transform arises as an iterated sigmoid transform in a stick-breaking
        construction of the `Dirichlet` distribution: the first logit is
        transformed via sigmoid to the first probability and the probability of
        everything else, and then the process recurses.

        This is bijective and appropriate for use in HMC; however it mixes
        coordinates together and is less appropriate for optimization.
        """

        domain = constraints.real_vector
        codomain = constraints.simplex
        bijective = True

        def __eq__(self, other):
            return isinstance(other, StickBreakingTransform)

        def _call(self, x):
            offset = x.shape[-1] + 1 - x.new_ones(x.shape[-1]).cumsum(-1)
            z = _clipped_sigmoid(x - offset.log())
            z_cumprod = (1 - z).cumprod(-1)
            y = pad(z, [0, 1], value=1) * pad(z_cumprod, [1, 0], value=1)
            return y

        def _inverse(self, y):
            y_crop = y[..., :-1]
            offset = y.shape[-1] - y.new_ones(y_crop.shape[-1]).cumsum(-1)
            sf = 1 - y_crop.cumsum(-1)
            # we clamp to make sure that sf is positive which sometimes does not
            # happen when y[-1] ~ 0 or y[:-1].sum() ~ 1
            sf = torch.clamp(sf, min=torch.finfo(y.dtype).tiny)
            x = y_crop.log() - sf.log() + offset.log()
            return x

        def log_abs_det_jacobian(self, x, y):
            offset = x.shape[-1] + 1 - x.new_ones(x.shape[-1]).cumsum(-1)
            x = x - offset.log()
            # use the identity 1 - sigmoid(x) = exp(-x) * sigmoid(x)
            detJ = (-x + F.logsigmoid(x) + y[..., :-1].log()).sum(-1)
            return detJ

        def forward_shape(self, shape):
            if len(shape) < 1:
                raise ValueError("Too few dimensions on input")
            return shape[:-1] + (shape[-1] + 1,)

        def inverse_shape(self, shape):
            if len(shape) < 1:
                raise ValueError("Too few dimensions on input")
            return shape[:-1] + (shape[-1] - 1,)

    class LowerCholeskyTransform(Transform):
        """
        Transform from unconstrained matrices to lower-triangular matrices with
        nonnegative diagonal entries.

        This is useful for parameterizing positive definite matrices in terms of
        their Cholesky factorization.
        """

        domain = constraints.independent(constraints.real, 2)
        codomain = constraints.lower_cholesky

        def __eq__(self, other):
            return isinstance(other, LowerCholeskyTransform)

        def _call(self, x):
            return x.tril(-1) + x.diagonal(dim1=-2, dim2=-1).exp().diag_embed()

        def _inverse(self, y):
            return y.tril(-1) + y.diagonal(dim1=-2, dim2=-1).log().diag_embed()

    class PositiveDefiniteTransform(Transform):
        """
        Transform from unconstrained matrices to positive-definite matrices.
        """

        domain = constraints.independent(constraints.real, 2)
        codomain = constraints.positive_definite  # type: ignore[assignment]

        def __eq__(self, other):
            return isinstance(other, PositiveDefiniteTransform)

        def _call(self, x):
            x = LowerCholeskyTransform()(x)
            return x @ x.mT

        def _inverse(self, y):
            y = torch.linalg.cholesky(y)
            return LowerCholeskyTransform().inv(y)

    class CatTransform(Transform):
        """
        Transform functor that applies a sequence of transforms `tseq`
        component-wise to each submatrix at `dim`, of length `lengths[dim]`,
        in a way compatible with :func:`torch.cat`.

        Example::

           x0 = torch.cat([torch.range(1, 10), torch.range(1, 10)], dim=0)
           x = torch.cat([x0, x0], dim=0)
           t0 = CatTransform([ExpTransform(), identity_transform], dim=0, lengths=[10, 10])
           t = CatTransform([t0, t0], dim=0, lengths=[20, 20])
           y = t(x)
        """

        transforms: List[Transform]

        def __init__(self, tseq, dim=0, lengths=None, cache_size=0):
            assert all(isinstance(t, Transform) for t in tseq)
            if cache_size:
                tseq = [t.with_cache(cache_size) for t in tseq]
            super().__init__(cache_size=cache_size)
            self.transforms = list(tseq)
            if lengths is None:
                lengths = [1] * len(self.transforms)
            self.lengths = list(lengths)
            assert len(self.lengths) == len(self.transforms)
            self.dim = dim

        @lazy_property
        def event_dim(self):
            return max(t.event_dim for t in self.transforms)

        @lazy_property
        def length(self):
            return sum(self.lengths)

        def with_cache(self, cache_size=1):
            if self._cache_size == cache_size:
                return self
            return CatTransform(self.transforms, self.dim, self.lengths, cache_size)

        def _call(self, x):
            assert -x.dim() <= self.dim < x.dim()
            assert x.size(self.dim) == self.length
            yslices = []
            start = 0
            for trans, length in zip(self.transforms, self.lengths):
                xslice = x.narrow(self.dim, start, length)
                yslices.append(trans(xslice))
                start = start + length  # avoid += for jit compat
            return torch.cat(yslices, dim=self.dim)

        def _inverse(self, y):
            assert -y.dim() <= self.dim < y.dim()
            assert y.size(self.dim) == self.length
            xslices = []
            start = 0
            for trans, length in zip(self.transforms, self.lengths):
                yslice = y.narrow(self.dim, start, length)
                xslices.append(trans.inv(yslice))
                start = start + length  # avoid += for jit compat
            return torch.cat(xslices, dim=self.dim)

        def log_abs_det_jacobian(self, x, y):
            assert -x.dim() <= self.dim < x.dim()
            assert x.size(self.dim) == self.length
            assert -y.dim() <= self.dim < y.dim()
            assert y.size(self.dim) == self.length
            logdetjacs = []
            start = 0
            for trans, length in zip(self.transforms, self.lengths):
                xslice = x.narrow(self.dim, start, length)
                yslice = y.narrow(self.dim, start, length)
                logdetjac = trans.log_abs_det_jacobian(xslice, yslice)
                if trans.event_dim < self.event_dim:
                    logdetjac = _sum_rightmost(logdetjac, self.event_dim - trans.event_dim)
                logdetjacs.append(logdetjac)
                start = start + length  # avoid += for jit compat
            # Decide whether to concatenate or sum.
            dim = self.dim
            if dim >= 0:
                dim = dim - x.dim()
            dim = dim + self.event_dim
            if dim < 0:
                return torch.cat(logdetjacs, dim=dim)
            else:
                return sum(logdetjacs)

        @property
        def bijective(self):
            return all(t.bijective for t in self.transforms)

        @constraints.dependent_property
        def domain(self):
            return constraints.cat(
                [t.domain for t in self.transforms], self.dim, self.lengths
            )

        @constraints.dependent_property
        def codomain(self):
            return constraints.cat(
                [t.codomain for t in self.transforms], self.dim, self.lengths
            )

    class StackTransform(Transform):
        """
        Transform functor that applies a sequence of transforms `tseq`
        component-wise to each submatrix at `dim`
        in a way compatible with :func:`torch.stack`.

        Example::

           x = torch.stack([torch.range(1, 10), torch.range(1, 10)], dim=1)
           t = StackTransform([ExpTransform(), identity_transform], dim=1)
           y = t(x)
        """

        transforms: List[Transform]

        def __init__(self, tseq, dim=0, cache_size=0):
            assert all(isinstance(t, Transform) for t in tseq)
            if cache_size:
                tseq = [t.with_cache(cache_size) for t in tseq]
            super().__init__(cache_size=cache_size)
            self.transforms = list(tseq)
            self.dim = dim

        def with_cache(self, cache_size=1):
            if self._cache_size == cache_size:
                return self
            return StackTransform(self.transforms, self.dim, cache_size)

        def _slice(self, z):
            return [z.select(self.dim, i) for i in range(z.size(self.dim))]

        def _call(self, x):
            assert -x.dim() <= self.dim < x.dim()
            assert x.size(self.dim) == len(self.transforms)
            yslices = []
            for xslice, trans in zip(self._slice(x), self.transforms):
                yslices.append(trans(xslice))
            return torch.stack(yslices, dim=self.dim)

        def _inverse(self, y):
            assert -y.dim() <= self.dim < y.dim()
            assert y.size(self.dim) == len(self.transforms)
            xslices = []
            for yslice, trans in zip(self._slice(y), self.transforms):
                xslices.append(trans.inv(yslice))
            return torch.stack(xslices, dim=self.dim)

        def log_abs_det_jacobian(self, x, y):
            assert -x.dim() <= self.dim < x.dim()
            assert x.size(self.dim) == len(self.transforms)
            assert -y.dim() <= self.dim < y.dim()
            assert y.size(self.dim) == len(self.transforms)
            logdetjacs = []
            yslices = self._slice(y)
            xslices = self._slice(x)
            for xslice, yslice, trans in zip(xslices, yslices, self.transforms):
                logdetjacs.append(trans.log_abs_det_jacobian(xslice, yslice))
            return torch.stack(logdetjacs, dim=self.dim)

        @property
        def bijective(self):
            return all(t.bijective for t in self.transforms)

        @constraints.dependent_property
        def domain(self):
            return constraints.stack([t.domain for t in self.transforms], self.dim)

        @constraints.dependent_property
        def codomain(self):
            return constraints.stack([t.codomain for t in self.transforms], self.dim)

    class CumulativeDistributionTransform(Transform):
        """
        Transform via the cumulative distribution function of a probability distribution.

        Args:
            distribution (Distribution): Distribution whose cumulative distribution function to use for
                the transformation.

        Example::

            # Construct a Gaussian copula from a multivariate normal.
            base_dist = MultivariateNormal(
                loc=torch.zeros(2),
                scale_tril=LKJCholesky(2).sample(),
            )
            transform = CumulativeDistributionTransform(Normal(0, 1))
            copula = TransformedDistribution(base_dist, [transform])
        """

        bijective = True
        codomain = constraints.unit_interval
        sign = +1

        def __init__(self, distribution, cache_size=0):
            super().__init__(cache_size=cache_size)
            self.distribution = distribution

        @property
        def domain(self):
            return self.distribution.support

        def _call(self, x):
            return self.distribution.cdf(x)

        def _inverse(self, y):
            return self.distribution.icdf(y)

        def log_abs_det_jacobian(self, x, y):
            return self.distribution.log_prob(x)

        def with_cache(self, cache_size=1):
            if self._cache_size == cache_size:
                return self
            return CumulativeDistributionTransform(self.distribution, cache_size=cache_size)

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.