zuko.transforms#
Parameterizable transformations.
Classes#
Creates a transformation \(f(x) = f_n \circ \dots \circ f_0(x)\). |
|
Wraps a base transformation to treat right-most dimensions as dependent. |
|
Creates a transformation \(f(x) = x\). |
|
Creates a transformation \(f(x) = -\cos(x)\). |
|
Creates a transformation \(f(x) = \sin(x)\). |
|
Creates a transformation that maps \(\mathbb{R}\) to the interval \([-B, B]\). |
|
Creates a transformation that circularly shifts the interval \([-B, B]\). |
|
Creates a transformation \(f(x) = \mathrm{sign}(x) |x|^{\exp(\alpha)}\). |
|
Creates a transformation \(f(x) = \exp(a) x + b\). |
|
Creates a monotonic rational-quadratic spline (RQS) transformation. |
|
Creates a transformation from a monotonic univariate function \(f_\phi(x)\). |
|
Creates a monotonic Bernstein polynomial transformation. |
|
Creates a gaussianization transformation. |
|
Creates a monotonic transformation \(f(x)\) by integrating a positive univariate function \(g(x)\). |
|
Creates a sum-of-squares (SOS) polynomial transformation. |
|
Transform via an autoregressive scheme. |
|
Transform via a coupling scheme. |
|
Creates a free-form Jacobian transformation. |
|
Creates a transformation that permutes the elements. |
|
Creates a rotation transformation \(f(x) = R x\). |
|
Creates a linear transformation \(f(x) = L U x\). |
Descriptions#
- class zuko.transforms.ComposedTransform(*transforms, **kwargs)#
Creates a transformation \(f(x) = f_n \circ \dots \circ f_0(x)\).
Optimized version of
torch.distributions.transforms.ComposeTransform.- Parameters:
transforms (Transform) – A sequence of transformations \(f_i\).
- class zuko.transforms.DependentTransform(base, reinterpreted, **kwargs)#
Wraps a base transformation to treat right-most dimensions as dependent.
Optimized version of
torch.distributions.transforms.IndependentTransform.
- class zuko.transforms.IdentityTransform(cache_size=0)#
Creates a transformation \(f(x) = x\).
- class zuko.transforms.CosTransform(cache_size=0)#
Creates a transformation \(f(x) = -\cos(x)\).
- class zuko.transforms.SinTransform(cache_size=0)#
Creates a transformation \(f(x) = \sin(x)\).
- class zuko.transforms.SoftclipTransform(bound=1.0, **kwargs)#
Creates a transformation that maps \(\mathbb{R}\) to the interval \([-B, B]\).
\[f(x) = \frac{x}{1 + \left| \frac{x}{B} \right|}\]- Parameters:
bound (float) – The codomain bound \(B\).
- class zuko.transforms.CircularShiftTransform(bound=1.0, **kwargs)#
Creates a transformation that circularly shifts the interval \([-B, B]\).
\[f(x) = (x \bmod 2B) - B\]Note
This transformation is only bijective over its domain \([-B, B]\) as \(f(x) = f(x + 2kB)\) for all \(k \in \mathbb{Z}\).
- Parameters:
bound (float) – The domain bound \(B\).
- class zuko.transforms.SignedPowerTransform(alpha, **kwargs)#
Creates a transformation \(f(x) = \mathrm{sign}(x) |x|^{\exp(\alpha)}\).
- Parameters:
alpha (Tensor) – The unconstrained exponent \(\alpha\), with shape \((*,)\).
- class zuko.transforms.MonotonicAffineTransform(shift, scale, slope=0.0001, **kwargs)#
Creates a transformation \(f(x) = \exp(a) x + b\).
- class zuko.transforms.MonotonicRQSTransform(widths, heights, derivatives, bound=5.0, slope=0.0001, **kwargs)#
Creates a monotonic rational-quadratic spline (RQS) transformation.
References
Neural Spline Flows (Durkan et al., 2019)- Parameters:
widths (Tensor) – The unconstrained bin widths, with shape \((*, K)\).
heights (Tensor) – The unconstrained bin heights, with shape \((*, K)\).
derivatives (Tensor) – The unconstrained knot derivatives, with shape \((*, K - 1)\).
bound (float) – The spline’s (co)domain bound \(B\).
slope (float) – The minimum slope of the transformation.
- class zuko.transforms.MonotonicTransform(f=None, phi=(), bound=10.0, eps=1e-06, **kwargs)#
Creates a transformation from a monotonic univariate function \(f_\phi(x)\).
The inverse function \(f_\phi^{-1}\) is approximated using the bisection method.
- Parameters:
f (Callable[[Tensor], Tensor]) – A monotonic univariate function \(f_\phi\). If
None,self.fis used instead.phi (Iterable[Tensor]) – The parameters \(\phi\) of \(f_\phi\). Providing the parameters is required to make the inverse transformation trainable.
bound (float) – The domain bound \(B\).
eps (float) – The absolute tolerance for the inverse transformation.
- class zuko.transforms.BernsteinTransform(theta, linear=False, **kwargs)#
Creates a monotonic Bernstein polynomial transformation.
\[f(x) = \frac{1}{M + 1} \sum_{i=0}^{M} b_{i+1,M-i+1}(\sigma(x)) \, \theta_i\]where \(b_{i,j}\) are the Bernstein basis polynomials and \(\sigma(x)\) is the sigmoid function.
References
Short-Term Density Forecasting of Low-Voltage Load using Bernstein-Polynomial Normalizing Flows (Arpogaus et al., 2022)Wikipedia
https://wikipedia.org/wiki/Bernstein_polynomial
- Parameters:
theta (Tensor) – The unconstrained polynomial coefficients \(\theta\), with shape \((*, M + 1)\).
linear (bool) – Whether to replace the sigmoid function with a linear mapping \(\frac{x + B}{2B}\). If
True, input features are assumed to be in \([-B, B]\). Failing to satisfy this constraint will result in NaNs.kwargs – Keyword arguments passed to
MonotonicTransform.
- class zuko.transforms.GaussianizationTransform(shift, scale, **kwargs)#
Creates a gaussianization transformation.
\[f(x) = \Phi^{-1} \left( \frac{1}{K} \sum_{i=1}^K \Phi(\exp(a_i) x + b_i) \right)\]where \(\Phi\) is the cumulative distribution function (CDF) of the standard normal \(\mathcal{N}(0, 1)\).
References
Gaussianization (Chen et al., 2000)- Parameters:
shift (Tensor) – The shift terms \(b\), with shape \((*, K)\).
scale (Tensor) – The unconstrained scale factors \(a\), with shape \((*, K)\).
kwargs – Keyword arguments passed to
MonotonicTransform.
- class zuko.transforms.UnconstrainedMonotonicTransform(g=None, C=0.0, n=32, **kwargs)#
Creates a monotonic transformation \(f(x)\) by integrating a positive univariate function \(g(x)\).
\[f(x) = \int_0^x g(u) ~ du + C\]The definite integral is estimated by a \(n\)-point Gauss-Legendre quadrature.
- class zuko.transforms.SOSPolynomialTransform(a, C, **kwargs)#
Creates a sum-of-squares (SOS) polynomial transformation.
The transformation \(f(x)\) is expressed as the primitive integral of the sum of \(K\) squared polynomials of degree \(L\).
\[f(x) = \int_0^x \frac{1}{K} \sum_{i = 1}^K \left( 1 + \sum_{j = 0}^L a_{i,j} ~ u^j \right)^2 ~ du + C\]References
Sum-of-Squares Polynomial Flow (Jaini et al., 2019)- Parameters:
a (Tensor) – The polynomial coefficients \(a\), with shape \((*, K, L + 1)\).
C (Tensor) – The integration constant \(C\).
kwargs – Keyword arguments passed to
UnconstrainedMonotonicTransform.
- class zuko.transforms.AutoregressiveTransform(meta, passes, **kwargs)#
Transform via an autoregressive scheme.
\[y_i = f(x_i | x_{<i})\]
- class zuko.transforms.CouplingTransform(meta, mask, **kwargs)#
Transform via a coupling scheme.
\[\begin{split}y_a & = x_a \\ y_b & = f(x_b | x_a)\end{split}\]
- class zuko.transforms.FreeFormJacobianTransform(f, t0=0.0, t1=1.0, phi=(), atol=1e-06, rtol=1e-05, exact=True, **kwargs)#
Creates a free-form Jacobian transformation.
The transformation is the integration of a system of first-order ordinary differential equations
\[x(t_1) = x_0 + \int_{t_0}^{t_1} f_\phi(t, x(t)) ~ dt .\]References
FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models (Grathwohl et al., 2018)- Parameters:
f (Callable[[Tensor, Tensor], Tensor]) – A system of first-order ODEs \(f_\phi\).
phi (Iterable[Tensor]) – The parameters \(\phi\) of \(f_\phi\).
atol (float) – The absolute integration tolerance.
rtol (float) – The relative integration tolerance.
exact (bool) – Whether the exact log-determinant of the Jacobian or an unbiased stochastic estimate thereof is calculated.
- class zuko.transforms.PermutationTransform(order, **kwargs)#
Creates a transformation that permutes the elements.
- Parameters:
order (LongTensor) – The permutation order, with shape \((*, D)\).