causalflows.flows¶

Wrappers for causal normalizing flows using standard architectures.

Classes¶

`CausalFlow`	Creates a lazy causal normalizing flow (initialized after passing the context).
`CausalMAF`	Creates a causal flow using a masked autoregressive flow (MAF) as base model.
`CausalNAF`	Creates a causal flow using a neural autoregressive flow (NAF) as base model.
`CausalNCSF`	Creates a causal flow using a neural circular spline flow (NCSF) as base model.
`CausalNSF`	Creates a causal flow using a neural spline flow (NSF) as the base model.
`CausalUNAF`	Creates a causal flow using an unconstrained neural autoregressive flow (UNAF) as base model.

Descriptions¶

class causalflows.flows.CausalFlow(transform, base)¶

Creates a lazy causal normalizing flow (initialized after passing the context).

See also

zuko.flows.core.Flow - Flow base class from Zuko.
causalflows.distributions.CausalNormalizingFlow - Class returned after the forward pass.

Parameters:

transform (Union[LazyTransform, Sequence[LazyTransform]]) – A lazy transformation or sequence of lazy transformations.
base (LazyDistribution) – A lazy distribution.

forward(context=None)¶

Parameters:: context (Tensor | None) – An input tensor representing the context of the (conditional) flow.
Returns:: A causal normalizing flow \(p(X | c)\).
Return type:: CausalNormalizingFlow

class causalflows.flows.CausalMAF(features, context=0, *args, order=None, adjacency=None, **kwargs)¶

Creates a causal flow using a masked autoregressive flow (MAF) as base model.

Parameters:

features (int) – The number of features.
context (int) – The number of context features.
order (LongTensor | None) – The causal order to follow by the flow. If used, then adjacency must be None.
adjacency (BoolTensor | None) – The causal graph to pass to the flow. If used, then order must be None.
kwargs – Keyword arguments passed to MaskedAutoregressiveTransform.

See also

MAF The equivalent non-causal counterpart from Zuko.

Example

>>> flow = CausalMAF(3, 4, order=torch.arange(3))
>>> flow
CausalMAF(
  (transform): 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): UnconditionalDistribution(DiagNormal(loc: torch.Size([3]), scale: torch.Size([3])))
)
>>> c = torch.randn(4)
>>> x = flow(c).sample()
>>> x
tensor([-1.9301, -0.1411, -0.7982])
>>> flow(c).log_prob(x)
tensor(-5.1800, grad_fn=<AddBackward0>)

References

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

https://arxiv.org/abs/1705.07057

class causalflows.flows.CausalNAF(features, context=0, signal=16, *args, order=None, adjacency=None, network=None, **kwargs)¶

Creates a causal flow using a neural autoregressive flow (NAF) as base model.

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.
order (LongTensor | None) – The causal order to follow by the flow. If used, then adjacency must be None.
adjacency (BoolTensor | None) – The causal graph to follow by the flow. If used, then order must be None.
network (Dict[str, Any] | None) – Keyword arguments passed to MNN.
kwargs – Keyword arguments passed to MaskedAutoregressiveTransform.

See also

NAF The equivalent non-causal counterpart from Zuko.

Warning

Invertibility is only guaranteed for features within the interval \([-10, 10]\). It is recommended to standardize features (zero mean, unit variance) before training.

References

Neural Autoregressive Flows (Huang et al., 2018)

https://arxiv.org/abs/1804.00779

class causalflows.flows.CausalNCSF(features, context=0, bins=8, **kwargs)¶

Creates a causal flow using a neural circular spline flow (NCSF) as base model.

Circular spline transformations are obtained by composing circular domain shifts with regular spline transformations. Features are assumed to lie in the half-open interval \([-\pi, \pi)\).

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.

See also

NCSF The equivalent non-causal counterpart from Zuko.

References

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

https://arxiv.org/abs/2002.02428

class causalflows.flows.CausalNSF(features, context=0, bins=8, **kwargs)¶

Creates a causal flow using a neural spline flow (NSF) as the base model.

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 CausalMAF.

See also

NSF The equivalent non-causal counterpart from Zuko.

Warning

Spline transformations are defined over the domain \([-5, 5]\). Any feature outside of this domain is not transformed. It is recommended to standardize features (zero mean, unit variance) before training.

References

Neural Spline Flows (Durkan et al., 2019)

https://arxiv.org/abs/1906.04032

class causalflows.flows.CausalUNAF(features, context=0, signal=16, *args, order=None, adjacency=None, network=None, **kwargs)¶

Creates a causal flow using an unconstrained neural autoregressive flow (UNAF) as base model.

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.
order (LongTensor | None) – The causal order to follow by the flow. If used, then adjacency must be None.
adjacency (BoolTensor | None) – The causal graph to follow by the flow. If used, then order must be None.
network (Dict[str, Any]) – Keyword arguments passed to UMNN.
kwargs – Keyword arguments passed to MaskedAutoregressiveTransform.

See also

UNAF The equivalent non-causal counterpart from Zuko.

Warning

Invertibility is only guaranteed for features within the interval \([-10, 10]\). It is recommended to standardize features (zero mean, unit variance) before training.

References

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

https://arxiv.org/abs/1908.05164