# Copyright (c) 2017-2019 Uber Technologies, Inc.
# SPDX-License-Identifier: Apache-2.0
from torch.distributions.transforms import * # noqa F403
# isort: split
from torch.distributions import biject_to, transform_to
from torch.distributions.transforms import (
ComposeTransform,
ExpTransform,
LowerCholeskyTransform,
)
from torch.distributions.transforms import __all__ as torch_transforms
from .. import constraints
from ..torch_transform import ComposeTransformModule
from .affine_autoregressive import (
AffineAutoregressive,
ConditionalAffineAutoregressive,
affine_autoregressive,
conditional_affine_autoregressive,
)
from .affine_coupling import (
AffineCoupling,
ConditionalAffineCoupling,
affine_coupling,
conditional_affine_coupling,
)
from .basic import ELUTransform, LeakyReLUTransform, elu, leaky_relu
from .batchnorm import BatchNorm, batchnorm
from .block_autoregressive import BlockAutoregressive, block_autoregressive
from .cholesky import (
CholeskyTransform,
CorrCholeskyTransform,
CorrLCholeskyTransform,
CorrMatrixCholeskyTransform,
)
from .discrete_cosine import DiscreteCosineTransform
from .generalized_channel_permute import (
ConditionalGeneralizedChannelPermute,
GeneralizedChannelPermute,
conditional_generalized_channel_permute,
generalized_channel_permute,
)
from .haar import HaarTransform
from .householder import (
ConditionalHouseholder,
Householder,
conditional_householder,
householder,
)
from .lower_cholesky_affine import LowerCholeskyAffine
from .matrix_exponential import (
ConditionalMatrixExponential,
MatrixExponential,
conditional_matrix_exponential,
matrix_exponential,
)
from .neural_autoregressive import (
ConditionalNeuralAutoregressive,
NeuralAutoregressive,
conditional_neural_autoregressive,
neural_autoregressive,
)
from .normalize import Normalize
from .ordered import OrderedTransform
from .permute import Permute, permute
from .planar import ConditionalPlanar, Planar, conditional_planar, planar
from .polynomial import Polynomial, polynomial
from .power import PositivePowerTransform
from .radial import ConditionalRadial, Radial, conditional_radial, radial
from .simplex_to_ordered import SimplexToOrderedTransform
from .softplus import SoftplusLowerCholeskyTransform, SoftplusTransform
from .spline import ConditionalSpline, Spline, conditional_spline, spline
from .spline_autoregressive import (
ConditionalSplineAutoregressive,
SplineAutoregressive,
conditional_spline_autoregressive,
spline_autoregressive,
)
from .spline_coupling import SplineCoupling, spline_coupling
from .sylvester import Sylvester, sylvester
from .unit_cholesky import UnitLowerCholeskyTransform
########################################
# register transforms
@transform_to.register(constraints.sphere)
def _transform_to_sphere(constraint):
return Normalize()
@biject_to.register(constraints.corr_matrix)
@transform_to.register(constraints.corr_matrix)
def _transform_to_corr_matrix(constraint):
return ComposeTransform(
[CorrCholeskyTransform(), CorrMatrixCholeskyTransform().inv]
)
@biject_to.register(constraints.ordered_vector)
@transform_to.register(constraints.ordered_vector)
def _transform_to_ordered_vector(constraint):
return OrderedTransform()
@biject_to.register(constraints.positive_ordered_vector)
@transform_to.register(constraints.positive_ordered_vector)
def _transform_to_positive_ordered_vector(constraint):
return ComposeTransform([OrderedTransform(), ExpTransform()])
# TODO: register biject_to when LowerCholeskyTransform is bijective
@transform_to.register(constraints.positive_definite)
def _transform_to_positive_definite(constraint):
return ComposeTransform([LowerCholeskyTransform(), CholeskyTransform().inv])
@biject_to.register(constraints.softplus_positive)
@transform_to.register(constraints.softplus_positive)
def _transform_to_softplus_positive(constraint):
return SoftplusTransform()
@transform_to.register(constraints.softplus_lower_cholesky)
def _transform_to_softplus_lower_cholesky(constraint):
return SoftplusLowerCholeskyTransform()
@transform_to.register(constraints.unit_lower_cholesky)
def _transform_to_unit_lower_cholesky(constraint):
return UnitLowerCholeskyTransform()
[docs]def iterated(repeats, base_fn, *args, **kwargs):
"""
Helper function to compose a sequence of bijective transforms with potentially
learnable parameters using :class:`~pyro.distributions.ComposeTransformModule`.
:param repeats: number of repeated transforms.
:param base_fn: function to construct the bijective transform.
:param args: arguments taken by `base_fn`.
:param kwargs: keyword arguments taken by `base_fn`.
:return: instance of :class:`~pyro.distributions.TransformModule`.
"""
assert isinstance(repeats, int) and repeats >= 1
return ComposeTransformModule([base_fn(*args, **kwargs) for _ in range(repeats)])
__all__ = [
"iterated",
"AffineAutoregressive",
"AffineCoupling",
"BatchNorm",
"BlockAutoregressive",
"CholeskyTransform",
"ComposeTransformModule",
"ConditionalAffineAutoregressive",
"ConditionalAffineCoupling",
"ConditionalGeneralizedChannelPermute",
"ConditionalHouseholder",
"ConditionalMatrixExponential",
"ConditionalNeuralAutoregressive",
"ConditionalPlanar",
"ConditionalRadial",
"ConditionalSpline",
"ConditionalSplineAutoregressive",
"CorrLCholeskyTransform",
"CorrMatrixCholeskyTransform",
"DiscreteCosineTransform",
"ELUTransform",
"GeneralizedChannelPermute",
"HaarTransform",
"Householder",
"LeakyReLUTransform",
"LowerCholeskyAffine",
"MatrixExponential",
"NeuralAutoregressive",
"Normalize",
"OrderedTransform",
"Permute",
"Planar",
"Polynomial",
"PositivePowerTransform",
"Radial",
"SimplexToOrderedTransform",
"SoftplusLowerCholeskyTransform",
"SoftplusTransform",
"Spline",
"SplineAutoregressive",
"SplineCoupling",
"Sylvester",
"affine_autoregressive",
"affine_coupling",
"batchnorm",
"block_autoregressive",
"conditional_affine_autoregressive",
"conditional_affine_coupling",
"conditional_generalized_channel_permute",
"conditional_householder",
"conditional_matrix_exponential",
"conditional_neural_autoregressive",
"conditional_planar",
"conditional_radial",
"conditional_spline",
"conditional_spline_autoregressive",
"elu",
"generalized_channel_permute",
"householder",
"leaky_relu",
"matrix_exponential",
"neural_autoregressive",
"permute",
"planar",
"polynomial",
"radial",
"spline",
"spline_autoregressive",
"spline_coupling",
"sylvester",
]
__all__.extend(torch_transforms)
del torch_transforms