Flows

CNeuralAffine

A specification for the coupling-neural affine flow, also known as the real-valued non-volume-preserving (real NVP) flow.

Definition

Given an input \(\mathbf{z}\) with \(D\) elements, we first partition them into \(\mathbf{z}_{1:d}\) and \(\mathbf{z}_{d+1:D}\). Then the transformation is given as:

\[\begin{aligned} T(\mathbf{z}_{1:d}) &= \mathbf{z}_{1:d} \\ T(\mathbf{z}_{d+1:D}) &= \mathbf{s}(\mathbf{z}_{1:d}) \odot \mathbf{z}_{d+1:D} + \mathbf{c}(\mathbf{z}_{1:d}) \end{aligned}\]

Where \(\mathbf{z} \in \mathbb{R}^D\), \(\mathbf{s} \in \mathbb{R}^{D/2}\) is non-negative, and \(\mathbf{c} \in \mathbb{R}^{D/2}\). Both \(\mathbf{s}\) and \(\mathbf{c}\) are joint outputs of a neural network.

Attributes:

Name	Type	Description
hidden_size	int	The width of the hidden layers for the neural network.
depth	int	The number of hidden layers for the neural network.
activation	Callable[[Array], Array]	The activation function to be used in between layers for the neural network.
flip	bool	Whether to flip the mask (which elements are the 'conditioner's and which are the 'transformer's).
key	PRNGKeyArray	The PRNG key used to generate the flow layer.

Source code in src/bayinx/flows/cneural_affine.py

class CNeuralAffine(FlowSpec):
    """
    A specification for the coupling-neural affine flow, also known as the real-valued non-volume-preserving (real NVP) flow.

    Definition:
        Given an input $\\mathbf{z}$ with $D$ elements, we first partition them into $\\mathbf{z}_{1:d}$ and $\\mathbf{z}_{d+1:D}$.
        Then the transformation is given as:

        $$\\begin{aligned}
        T(\\mathbf{z}_{1:d}) &= \\mathbf{z}_{1:d} \\\\
        T(\\mathbf{z}_{d+1:D}) &= \\mathbf{s}(\\mathbf{z}_{1:d}) \\odot \\mathbf{z}_{d+1:D} + \\mathbf{c}(\\mathbf{z}_{1:d})
        \\end{aligned}$$

        Where $\\mathbf{z} \\in \\mathbb{R}^D$, $\\mathbf{s} \\in \\mathbb{R}^{D/2}$ is non-negative, and $\\mathbf{c} \\in \\mathbb{R}^{D/2}$.
        Both $\\mathbf{s}$ and $\\mathbf{c}$ are joint outputs of a neural network.

    Attributes:
        hidden_size: The width of the hidden layers for the neural network.
        depth: The number of hidden layers for the neural network.
        activation: The activation function to be used in between layers for the neural network.
        flip: Whether to flip the mask (which elements are the 'conditioner's and which are the 'transformer's).
        key: The PRNG key used to generate the flow layer.
    """
    hidden_size: int
    depth: int
    activation: Callable[[Array], Array]
    flip: bool
    key: PRNGKeyArray

    def __init__(self, hidden_size: int = 16, depth: int = 2, activation: Callable[[Array], Array] = jnn.relu, flip: bool = False, key: PRNGKeyArray = jr.key(0)):
        self.hidden_size = hidden_size
        self.depth = depth
        self.activation = activation
        self.flip = flip
        self.key = key

    def construct(self, dim: int) -> CNeuralAffineLayer:
        return CNeuralAffineLayer(dim, self.hidden_size, self.depth, self.activation, self.flip, self.key)

DiagAffine

A specification for the diagonal affine flow.

Definition

\(T(\mathbf{z}) = \mathbf{d} \odot \mathbf{z} + \mathbf{c}\)

Where \(\mathbf{z} \in \mathbb{R}^D\), \(\mathbf{d} \in \mathbb{R}^{D}\) is non-negative, and \(\mathbf{c} \in \mathbb{R}^D\).

Attributes:

Name	Type	Description
key	PRNGKeyArray	The PRNG key used to generate the diagonal affine flow layer.

Source code in src/bayinx/flows/diagaffine.py

class DiagAffine(FlowSpec):
    """
    A specification for the diagonal affine flow.

    Definition:
        $T(\\mathbf{z}) = \\mathbf{d} \\odot \\mathbf{z} + \\mathbf{c}$

        Where $\\mathbf{z} \\in \\mathbb{R}^D$, $\\mathbf{d} \\in \\mathbb{R}^{D}$ is non-negative, and $\\mathbf{c} \\in \\mathbb{R}^D$.

    Attributes:
        key: The PRNG key used to generate the diagonal affine flow layer.
    """
    key: PRNGKeyArray
    def __init__(self, key: PRNGKeyArray = jr.key(0)):
        self.key = key

    def construct(self, dim: int) -> DiagAffineLayer:
        return DiagAffineLayer(dim, self.key)

FullAffine

A specification for the full affine flow.

Definition

\(T(\mathbf{z}) = \mathbf{L z} + \mathbf{c}\)

Where \(\mathbf{z} \in \mathbb{R}^D\), \(\mathbf{L} \in \mathbb{R}^{D, D}\) is lower triangular with a non-negative diagonal, and \(\mathbf{c} \in \mathbb{R}^D\).

Attributes:

Name	Type	Description
key	PRNGKeyArray	The PRNG key used to generate the full affine flow layer.

Source code in src/bayinx/flows/fullaffine.py

class FullAffine(FlowSpec):
    """
    A specification for the full affine flow.

    Definition:
        $T(\\mathbf{z}) = \\mathbf{L z} + \\mathbf{c}$

        Where $\\mathbf{z} \\in \\mathbb{R}^D$, $\\mathbf{L} \\in \\mathbb{R}^{D, D}$ is lower triangular with a non-negative diagonal, and $\\mathbf{c} \\in \\mathbb{R}^D$.

    Attributes:
        key: The PRNG key used to generate the full affine flow layer.
    """
    key: PRNGKeyArray

    def __init__(self, key: PRNGKeyArray = jr.key(0)):
        """
        Initializes the specification for a full affine flow.

        Parameters:
            key: A PRNG key used to generate the full affine flow.
        """
        self.key = key

    def construct(self, dim: int) -> FullAffineLayer:
        """
        Constructs a full affine flow layer.

        Parameters:
            dim: The dimension of the parameter space.

        Returns:
            A FullAffineLayer of dimension `dim`.
        """
        return FullAffineLayer(dim, self.key)

__init__(key: PRNGKeyArray = jr.key(0))

Initializes the specification for a full affine flow.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	A PRNG key used to generate the full affine flow.	key(0)

Source code in src/bayinx/flows/fullaffine.py

def __init__(self, key: PRNGKeyArray = jr.key(0)):
    """
    Initializes the specification for a full affine flow.

    Parameters:
        key: A PRNG key used to generate the full affine flow.
    """
    self.key = key

construct(dim: int) -> FullAffineLayer

Constructs a full affine flow layer.

Parameters:

Name	Type	Description	Default
dim	int	The dimension of the parameter space.	required

Returns:

Type	Description
FullAffineLayer	A FullAffineLayer of dimension dim.

Source code in src/bayinx/flows/fullaffine.py

def construct(self, dim: int) -> FullAffineLayer:
    """
    Constructs a full affine flow layer.

    Parameters:
        dim: The dimension of the parameter space.

    Returns:
        A FullAffineLayer of dimension `dim`.
    """
    return FullAffineLayer(dim, self.key)

LinearRationalSpline

A specification for a modified linear rational splines flow.

Definition

The transformation \(T(z)\) is defined as:

\[ T(\tilde{z}) = \begin{cases} y_0 + \delta_0 (\tilde{z} - x_0)(1 + |\tilde{z} - x_0|)^\alpha & z < x_0 \\ \text{LR}_b(\tilde{z}) & x_b \leq \tilde{z} < x_{b+1} \\ \tilde{z} & \tilde{z} > x_K \end{cases} \]

Where for each dimension \(d \in \{1, \dots, D\}\) we centre the spline at \(\tilde{z} = z - s\), the knots are denoted by \(\{x_k, y_k\}_{k=0}^K\) and derivatives by \(\{\delta_k\}_{k=0}^K\), with shift \(s\). The linear-rational segment for bin \(b\) is denoted by \(\text{LR}_b(z)\), and is parameterized by its knots \((x_b, y_b), (x_{b+1}, y_{b+1})\), derivatives at each knot \(\delta_b, \delta_{b+1}\), and midpoint weight \(\lambda_b\) controlling the curvature.

Attributes:

Name	Type	Description
n_bins	int	The number of bins covering the boundary.
key	PRNGKeyArray	The PRNG key used to generate the diagonal affine flow layer.

Source code in src/bayinx/flows/lrs.py

class LinearRationalSpline(FlowSpec):
    """
    A specification for a modified linear rational splines flow.

    Definition:
        The transformation $T(z)$ is defined as:

        $$
            T(\\tilde{z}) = \\begin{cases}
            y_0 + \\delta_0 (\\tilde{z} - x_0)(1 + |\\tilde{z} - x_0|)^\\alpha & z < x_0 \\\\
            \\text{LR}_b(\\tilde{z}) & x_b \\leq \\tilde{z} < x_{b+1} \\\\
            \\tilde{z} & \\tilde{z} > x_K
            \\end{cases}
        $$

        Where for each dimension $d \\in \\{1, \\dots, D\\}$ we centre the
        spline at $\\tilde{z} = z - s$, the knots are denoted by
        $\\{x_k, y_k\\}_{k=0}^K$ and derivatives by $\\{\\delta_k\\}_{k=0}^K$,
        with shift $s$. The linear-rational segment for bin $b$ is denoted by
        $\\text{LR}_b(z)$, and is parameterized by its knots
        $(x_b, y_b), (x_{b+1}, y_{b+1})$, derivatives at each knot
        $\\delta_b, \\delta_{b+1}$, and midpoint weight $\\lambda_b$
        controlling the curvature.

    Attributes:
        n_bins: The number of bins covering the boundary.
        key: The PRNG key used to generate the diagonal affine flow layer.
    """
    n_bins: int
    key: PRNGKeyArray

    def __init__(
        self,
        n_bins: int = 7,
        key: PRNGKeyArray = jr.key(0)
    ):
        self.n_bins = n_bins
        self.key = key

    def construct(self, dim: int) -> LinearRationalSplineLayer:
        return LinearRationalSplineLayer(dim, self.n_bins, self.key)

RealNVP

A specification for the coupling-neural affine flow, also known as the real-valued non-volume-preserving (real NVP) flow.

Definition

Given an input \(\mathbf{z}\) with \(D\) elements, we first partition them into \(\mathbf{z}_{1:d}\) and \(\mathbf{z}_{d+1:D}\). Then the transformation is given as:

\[\begin{aligned} T(\mathbf{z}_{1:d}) &= \mathbf{z}_{1:d} \\ T(\mathbf{z}_{d+1:D}) &= \mathbf{s}(\mathbf{z}_{1:d}) \odot \mathbf{z}_{d+1:D} + \mathbf{c}(\mathbf{z}_{1:d}) \end{aligned}\]

Where \(\mathbf{z} \in \mathbb{R}^D\), \(\mathbf{s} \in \mathbb{R}^{D/2}\) is non-negative, and \(\mathbf{c} \in \mathbb{R}^{D/2}\). Both \(\mathbf{s}\) and \(\mathbf{c}\) are joint outputs of a neural network.

Attributes:

Name	Type	Description
hidden_size	int	The width of the hidden layers for the neural network.
depth	int	The number of hidden layers for the neural network.
activation	Callable[[Array], Array]	The activation function to be used in between layers for the neural network.
flip	bool	Whether to flip the mask (which elements are the 'conditioner's and which are the 'transformer's).
key	PRNGKeyArray	The PRNG key used to generate the flow layer.

Source code in src/bayinx/flows/cneural_affine.py

class CNeuralAffine(FlowSpec):
    """
    A specification for the coupling-neural affine flow, also known as the real-valued non-volume-preserving (real NVP) flow.

    Definition:
        Given an input $\\mathbf{z}$ with $D$ elements, we first partition them into $\\mathbf{z}_{1:d}$ and $\\mathbf{z}_{d+1:D}$.
        Then the transformation is given as:

        $$\\begin{aligned}
        T(\\mathbf{z}_{1:d}) &= \\mathbf{z}_{1:d} \\\\
        T(\\mathbf{z}_{d+1:D}) &= \\mathbf{s}(\\mathbf{z}_{1:d}) \\odot \\mathbf{z}_{d+1:D} + \\mathbf{c}(\\mathbf{z}_{1:d})
        \\end{aligned}$$

        Where $\\mathbf{z} \\in \\mathbb{R}^D$, $\\mathbf{s} \\in \\mathbb{R}^{D/2}$ is non-negative, and $\\mathbf{c} \\in \\mathbb{R}^{D/2}$.
        Both $\\mathbf{s}$ and $\\mathbf{c}$ are joint outputs of a neural network.

    Attributes:
        hidden_size: The width of the hidden layers for the neural network.
        depth: The number of hidden layers for the neural network.
        activation: The activation function to be used in between layers for the neural network.
        flip: Whether to flip the mask (which elements are the 'conditioner's and which are the 'transformer's).
        key: The PRNG key used to generate the flow layer.
    """
    hidden_size: int
    depth: int
    activation: Callable[[Array], Array]
    flip: bool
    key: PRNGKeyArray

    def __init__(self, hidden_size: int = 16, depth: int = 2, activation: Callable[[Array], Array] = jnn.relu, flip: bool = False, key: PRNGKeyArray = jr.key(0)):
        self.hidden_size = hidden_size
        self.depth = depth
        self.activation = activation
        self.flip = flip
        self.key = key

    def construct(self, dim: int) -> CNeuralAffineLayer:
        return CNeuralAffineLayer(dim, self.hidden_size, self.depth, self.activation, self.flip, self.key)