zuko.flows#

Parameterized flows and autoregressive transformations.

Classes#

`DistributionModule`	Abstract distribution module.
`TransformModule`	Abstract transformation module.
`FlowModule`	Creates a normalizing flow module.
`GMM`	Creates a Gaussian mixture model (GMM).
`MaskedAutoregressiveTransform`	Creates a masked autoregressive transformation.
`MAF`	Creates a masked autoregressive flow (MAF).
`NSF`	Creates a neural spline flow (NSF) with monotonic rational-quadratic spline transformations.
`NCSF`	Creates a neural circular spline flow (NCSF).
`SOSPF`	Creates a sum-of-squares polynomial flow (SOSPF).
`NeuralAutoregressiveTransform`	Creates a neural autoregressive transformation.
`NAF`	Creates a neural autoregressive flow (NAF).
`UnconstrainedNeuralAutoregressiveTransform`	Creates an unconstrained neural autoregressive transformation.
`UNAF`	Creates an unconstrained neural autoregressive flow (UNAF).
`FFJTransform`	Creates a free-form Jacobian (FFJ) transformation.
`CNF`	Creates a continuous normalizing flow (CNF) with free-form Jacobian transformations.

Descriptions#

class zuko.flows.DistributionModule(*args, **kwargs)#

Abstract distribution module.

abstract forward()#

Parameters: y (Optional[Tensor]) – A context \(y\).
Returns: A distribution \(p(X | y)\).
Return type: Distribution

class zuko.flows.TransformModule(*args, **kwargs)#

Abstract transformation module.

abstract forward()#

Parameters: y (Optional[Tensor]) – A context \(y\).
Returns: A transformation \(z = f(x | y)\).
Return type: Transform

class zuko.flows.FlowModule(transforms, base)#

Creates a normalizing flow module.

Parameters

transforms (Sequence[TransformModule]) – A list of transformation modules.
base (DistributionModule) – A distribution module.

forward(y=None)#

Parameters: y (Optional[Tensor]) – A context \(y\).
Returns: A normalizing flow \(p(X | y)\).
Return type: NormalizingFlow

class zuko.flows.GMM(features, context=0, components=2, **kwargs)#

Creates a Gaussian mixture model (GMM).

\[p(X | y) = \sum_{i = 1}^K w_i(y) \, \mathcal{N}(X | \mu_i(y), \Sigma_i(y))\]

Parameters

features (int) – The number of features.
context (int) – The number of context features.
components (int) – The number of components \(K\) in the mixture.
**kwargs – Keyword arguments passed to zuko.nn.MLP.

class zuko.flows.MaskedAutoregressiveTransform(features, context=0, passes=None, order=None, univariate=<class 'zuko.transforms.MonotonicAffineTransform'>, shapes=((), ()), **kwargs)#

Creates a masked autoregressive transformation.

References

Masked Autoregressive Flow for Density Estimation (Papamakarios et al., 2017)

https://arxiv.org/abs/1705.07057

Parameters

features (int) – The number of features.
context (int) – The number of context features.
passes (int) – The number of sequential passes for the inverse transformation. If None, use the number of features instead, making the transformation fully autoregressive. Coupling corresponds to passes=2.
order (LongTensor) – The feature ordering. If None, use range(features) instead.
univariate (Callable[[...], Transform]) – The univariate transformation constructor.
shapes (Sequence[Size]) – The shapes of the univariate transformation parameters.
kwargs – Keyword arguments passed to zuko.nn.MaskedMLP.

Example

>>> t = MaskedAutoregressiveTransform(3, 4)
>>> t
MaskedAutoregressiveTransform(
  (base): MonotonicAffineTransform()
  (order): [0, 1, 2]
  (hyper): MaskedMLP(
    (0): MaskedLinear(in_features=7, out_features=64, bias=True)
    (1): ReLU()
    (2): MaskedLinear(in_features=64, out_features=64, bias=True)
    (3): ReLU()
    (4): MaskedLinear(in_features=64, out_features=6, bias=True)
  )
)
>>> x = torch.randn(3)
>>> x
tensor([-0.9485,  1.5290,  0.2018])
>>> y = torch.randn(4)
>>> z = t(y)(x)
>>> t(y).inv(z)
tensor([-0.9485,  1.5290,  0.2018])

class zuko.flows.MAF(features, context=0, transforms=3, randperm=False, **kwargs)#

Creates a masked autoregressive flow (MAF).

References

Masked Autoregressive Flow for Density Estimation (Papamakarios et al., 2017)

https://arxiv.org/abs/1705.07057

Parameters

features (int) – The number of features.
context (int) – The number of context features.
transforms (int) – The number of autoregressive transformations.
randperm (bool) – Whether features are randomly permuted between transformations or not. If False, features are in ascending (descending) order for even (odd) transformations.
kwargs – Keyword arguments passed to MaskedAutoregressiveTransform.

Example

>>> flow = MAF(3, 4, transforms=3)
>>> flow
MAF(
  (transforms): ModuleList(
    (0): MaskedAutoregressiveTransform(
      (base): MonotonicAffineTransform()
      (order): [0, 1, 2]
      (hyper): MaskedMLP(
        (0): MaskedLinear(in_features=7, out_features=64, bias=True)
        (1): ReLU()
        (2): MaskedLinear(in_features=64, out_features=64, bias=True)
        (3): ReLU()
        (4): MaskedLinear(in_features=64, out_features=6, bias=True)
      )
    )
    (1): MaskedAutoregressiveTransform(
      (base): MonotonicAffineTransform()
      (order): [2, 1, 0]
      (hyper): MaskedMLP(
        (0): MaskedLinear(in_features=7, out_features=64, bias=True)
        (1): ReLU()
        (2): MaskedLinear(in_features=64, out_features=64, bias=True)
        (3): ReLU()
        (4): MaskedLinear(in_features=64, out_features=6, bias=True)
      )
    )
    (2): MaskedAutoregressiveTransform(
      (base): MonotonicAffineTransform()
      (order): [0, 1, 2]
      (hyper): MaskedMLP(
        (0): MaskedLinear(in_features=7, out_features=64, bias=True)
        (1): ReLU()
        (2): MaskedLinear(in_features=64, out_features=64, bias=True)
        (3): ReLU()
        (4): MaskedLinear(in_features=64, out_features=6, bias=True)
      )
    )
  )
  (base): DiagNormal(loc: torch.Size([3]), scale: torch.Size([3]))
)
>>> y = torch.randn(4)
>>> x = flow(y).sample()
>>> x
tensor([-1.7154, -0.4401,  0.7505])
>>> flow(y).log_prob(x)
tensor(-4.4630, grad_fn=<AddBackward0>)

class zuko.flows.NSF(features, context=0, bins=8, **kwargs)#

Creates a neural spline flow (NSF) with monotonic rational-quadratic spline transformations.

Note

By default, transformations are fully autoregressive. Coupling transformations can be obtained by setting passes=2.

References

Neural Spline Flows (Durkan et al., 2019)

https://arxiv.org/abs/1906.04032

Parameters

features (int) – The number of features.
context (int) – The number of context features.
bins (int) – The number of bins \(K\).
kwargs – Keyword arguments passed to MAF.

class zuko.flows.NCSF(features, context=0, **kwargs)#

Creates a neural circular spline flow (NCSF).

Note

Features are assumed to lie in the half-open interval \([-\pi, \pi[\).

References

Normalizing Flows on Tori and Spheres (Rezende et al., 2020)

https://arxiv.org/abs/2002.02428

Parameters

features (int) – The number of features.
context (int) – The number of context features.
kwargs – Keyword arguments passed to NSF.

class zuko.flows.SOSPF(features, context=0, degree=3, polynomials=2, **kwargs)#

Creates a sum-of-squares polynomial flow (SOSPF).

References

Sum-of-Squares Polynomial Flow (Jaini et al., 2019)

https://arxiv.org/abs/1905.02325

Parameters

features (int) – The number of features.
context (int) – The number of context features.
degree (int) – The degree \(L\) of polynomials.
polynomials (int) – The number of polynomials \(K\).
kwargs – Keyword arguments passed to MAF.

class zuko.flows.NeuralAutoregressiveTransform(features, context=0, signal=8, network={}, **kwargs)#

Creates a neural autoregressive transformation.

The monotonic neural network is parametrized by its internal positive weights, which are independent of the features and context. To modulate its behavior, it receives as input a signal that is autoregressively dependent on the features and context.

References

Neural Autoregressive Flows (Huang et al., 2018)

https://arxiv.org/abs/1804.00779

Parameters

features (int) – The number of features.
context (int) – The number of context features.
signal (int) – The number of signal features of the monotonic network.
network (Dict[str, Any]) – Keyword arguments passed to zuko.nn.MonotonicMLP.
kwargs – Keyword arguments passed to MaskedAutoregressiveTransform.

Example

>>> t = NeuralAutoregressiveTransform(3, 4)
>>> t
NeuralAutoregressiveTransform(
  (base): MonotonicTransform()
  (order): [0, 1, 2]
  (hyper): MaskedMLP(
    (0): MaskedLinear(in_features=7, out_features=64, bias=True)
    (1): ReLU()
    (2): MaskedLinear(in_features=64, out_features=64, bias=True)
    (3): ReLU()
    (4): MaskedLinear(in_features=64, out_features=24, bias=True)
  )
  (network): MonotonicMLP(
    (0): MonotonicLinear(in_features=9, out_features=64, bias=True)
    (1): TwoWayELU(alpha=1.0)
    (2): MonotonicLinear(in_features=64, out_features=64, bias=True)
    (3): TwoWayELU(alpha=1.0)
    (4): MonotonicLinear(in_features=64, out_features=1, bias=True)
  )
)
>>> x = torch.randn(3)
>>> x
tensor([-2.3267,  1.4581, -1.6776])
>>> y = torch.randn(4)
>>> z = t(y)(x)
>>> t(y).inv(z)
tensor([-2.3267,  1.4581, -1.6776])

class zuko.flows.NAF(features, context=0, transforms=3, randperm=False, **kwargs)#

Creates a neural autoregressive flow (NAF).

References

Neural Autoregressive Flows (Huang et al., 2018)

https://arxiv.org/abs/1804.00779

Parameters

features (int) – The number of features.
context (int) – The number of context features.
transforms (int) – The number of autoregressive transformations.
randperm (bool) – Whether features are randomly permuted between transformations or not. If False, features are in ascending (descending) order for even (odd) transformations.
unconstrained – Whether to use unconstrained or regular monotonic networks.
kwargs – Keyword arguments passed to NeuralAutoregressiveTransform.

class zuko.flows.UnconstrainedNeuralAutoregressiveTransform(features, context=0, signal=8, network={}, **kwargs)#

Creates an unconstrained neural autoregressive transformation.

The integrand neural network is parametrized by its internal weights, which are independent of the features and context. To modulate its behavior, it receives as input a signal that is autoregressively dependent on the features and context. The integration constant has the same dependencies as the signal.

References

Unconstrained Monotonic Neural Networks (Wehenkel et al., 2019)

https://arxiv.org/abs/1908.05164

Parameters

features (int) – The number of features.
context (int) – The number of context features.
signal (int) – The number of signal features of the integrand network.
network (Dict[str, Any]) – Keyword arguments passed to zuko.nn.MLP.
kwargs – Keyword arguments passed to MaskedAutoregressiveTransform.

Example

>>> t = UnconstrainedNeuralAutoregressiveTransform(3, 4)
>>> t
UnconstrainedNeuralAutoregressiveTransform(
  (base): UnconstrainedMonotonicTransform()
  (order): [0, 1, 2]
  (hyper): MaskedMLP(
    (0): MaskedLinear(in_features=7, out_features=64, bias=True)
    (1): ReLU()
    (2): MaskedLinear(in_features=64, out_features=64, bias=True)
    (3): ReLU()
    (4): MaskedLinear(in_features=64, out_features=27, bias=True)
  )
  (integrand): MLP(
    (0): Linear(in_features=9, out_features=64, bias=True)
    (1): ELU(alpha=1.0)
    (2): Linear(in_features=64, out_features=64, bias=True)
    (3): ELU(alpha=1.0)
    (4): Linear(in_features=64, out_features=1, bias=True)
    (5): Softplus(beta=1, threshold=20)
  )
)
>>> x = torch.randn(3)
>>> x
tensor([-0.0103, -1.0871, -0.0667])
>>> y = torch.randn(4)
>>> z = t(y)(x)
>>> t(y).inv(z)
tensor([-0.0103, -1.0871, -0.0667])

class zuko.flows.UNAF(features, context=0, transforms=3, randperm=False, **kwargs)#

Creates an unconstrained neural autoregressive flow (UNAF).

References

Unconstrained Monotonic Neural Networks (Wehenkel et al., 2019)

https://arxiv.org/abs/1908.05164

Parameters

features (int) – The number of features.
context (int) – The number of context features.
transforms (int) – The number of autoregressive transformations.
randperm (bool) – Whether features are randomly permuted between transformations or not. If False, features are in ascending (descending) order for even (odd) transformations.
kwargs – Keyword arguments passed to UnconstrainedNeuralAutoregressiveTransform.

class zuko.flows.FFJTransform(features, context=0, frequencies=3, exact=True, **kwargs)#

Creates a free-form Jacobian (FFJ) transformation.

References

FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models (Grathwohl et al., 2018)

https://arxiv.org/abs/1810.01367

Parameters

features (int) – The number of features.
context (int) – The number of context features.
frequencies (int) – The number of time embedding frequencies.
exact (bool) – Whether the exact log-determinant of the Jacobian or an unbiased stochastic estimate thereof is calculated.
kwargs – Keyword arguments passed to zuko.nn.MLP.

Example

>>> t = FFJTransform(3, 4)
>>> t
FFJTransform(
  (ode): MLP(
    (0): Linear(in_features=13, out_features=64, bias=True)
    (1): ELU(alpha=1.0)
    (2): Linear(in_features=64, out_features=64, bias=True)
    (3): ELU(alpha=1.0)
    (4): Linear(in_features=64, out_features=3, bias=True)
  )
)
>>> x = torch.randn(3)
>>> x
tensor([ 0.1777,  1.0139, -1.0370])
>>> y = torch.randn(4)
>>> z = t(y)(x)
>>> t(y).inv(z)
tensor([ 0.1777,  1.0139, -1.0370])

class zuko.flows.CNF(features, context=0, transforms=1, **kwargs)#

Creates a continuous normalizing flow (CNF) with free-form Jacobian transformations.

References

Neural Ordinary Differential Equations (Chen el al., 2018)

https://arxiv.org/abs/1806.07366

FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models (Grathwohl et al., 2018)

https://arxiv.org/abs/1810.01367

Parameters

features (int) – The number of features.
context (int) – The number of context features.
transforms (int) – The number of transformations.
kwargs – Keyword arguments passed to FFJTransform.