zuko.flows.continuous#

Continuous flows and transformations.

Classes#

`FFJTransform`	Creates a lazy free-form Jacobian (FFJ) transformation.
`CNF`	Creates a continuous normalizing flow (CNF) with a free-form Jacobian transformation.

Descriptions#

class zuko.flows.continuous.FFJTransform(features, context=0, freqs=3, atol=1e-06, rtol=1e-05, exact=True, **kwargs)#

Creates a lazy 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.
freqs (int) – The number of time embedding frequencies.
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.
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])
>>> c = torch.randn(4)
>>> y = t(c)(x)
>>> t(c).inv(y)
tensor([ 0.1777,  1.0139, -1.0370])

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

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

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.
kwargs – Keyword arguments passed to FFJTransform.