Skip to content

Flows

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 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
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)):
        """
        Initializes the specification for a diagonal affine flow.

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

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

        Parameters:
            dim: The dimension of the parameter space.

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

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

Initializes the specification for a diagonal affine flow.

Parameters:

Name Type Description Default
key PRNGKeyArray

A PRNG key used to generate the diagonal affine flow.

 key(0)
Source code in src/bayinx/flows/diagaffine.py 
107
108
109
110
111
112
113
114
def __init__(self, key: PRNGKeyArray = jr.key(0)):
    """
    Initializes the specification for a diagonal affine flow.

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

construct(dim: int) -> DiagAffineLayer

Constructs a diagonal affine flow layer.

Parameters:

Name Type Description Default
dim int

The dimension of the parameter space.

 required

Returns:

Type Description
DiagAffineLayer

A DiagonalAffineLayer of dimension dim.
Source code in src/bayinx/flows/diagaffine.py 
116
117
118
119
120
121
122
123
124
125
126
def construct(self, dim: int) -> DiagAffineLayer:
    """
    Constructs a diagonal affine flow layer.

    Parameters:
        dim: The dimension of the parameter space.

    Returns:
        A DiagonalAffineLayer of dimension `dim`.
    """
    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 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
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 
129
130
131
132
133
134
135
136
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 
138
139
140
141
142
143
144
145
146
147
148
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)

LowRankAffine

A specification for the low-rank affine flow.

Definition

\(T(\mathbf{z}) = \mathbf{D z} + ((\mathbf{U V}^\top) \odot \mathbf{M})\mathbf{z} + \mathbf{c}\)

Where \(\mathbf{z} \in \mathbb{R}^D\), \(\mathbf{D} \in \mathbb{R}^{D, D}\) is a non-negative diagonal matrix, \(\mathbf{U}, \mathbf{V} \in \mathbb{R}^{D, R}\) are low-rank factor matrices, \(\mathbf{M} \in \mathbb{R}^{D, D}\) is an implicit strictly lower-triangular mask, and \(\mathbf{c} \in \mathbb{R}^D\) is the shift vector.

Attributes:

Name Type Description
rank int

The rank \(R\) of the low-rank factor matrices \(\mathbf{U}\) and \(\mathbf{V}\).
Source code in src/bayinx/flows/lowrankaffine.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
class LowRankAffine(FlowSpec):
    """
    A specification for the low-rank affine flow.

    Definition:
        $T(\\mathbf{z}) = \\mathbf{D z} + ((\\mathbf{U V}^\\top) \\odot \\mathbf{M})\\mathbf{z} + \\mathbf{c}$

        Where $\\mathbf{z} \\in \\mathbb{R}^D$, $\\mathbf{D} \\in \\mathbb{R}^{D, D}$ is a non-negative diagonal matrix,
        $\\mathbf{U}, \\mathbf{V} \\in \\mathbb{R}^{D, R}$ are low-rank factor matrices, $\\mathbf{M} \\in \\mathbb{R}^{D, D}$
        is an implicit strictly lower-triangular mask, and $\\mathbf{c} \\in \\mathbb{R}^D$ is the shift vector.

    Attributes:
        rank (int): The rank $R$ of the low-rank factor matrices $\\mathbf{U}$ and $\\mathbf{V}$.
    """
    rank: int

    def __init__(self, rank: int):
        """
        Initializes the specification for a low-rank affine flow.

        Parameters:
            rank: The rank R (number of columns in U and V) of the low-rank off-diagonal component.
        """
        self.rank = rank

    def construct(self, dim: int) -> LowRankAffineLayer:
        """
        Constructs a low-rank affine flow layer.

        Parameters:
            dim: The dimension of the parameter space.

        Returns:
            A LowRankAffineLayer of dimension `dim` and rank `self.rank`.
        """
        if (self.rank > (dim - 1)/2):
            raise ValueError(f"Rank {self.rank} is large, consider using a full affine flow instead.")

        return LowRankAffineLayer(dim, self.rank)

__init__(rank: int)

Initializes the specification for a low-rank affine flow.

Parameters:

Name Type Description Default
rank int

The rank R (number of columns in U and V) of the low-rank off-diagonal component.

 required
Source code in src/bayinx/flows/lowrankaffine.py 
171
172
173
174
175
176
177
178
def __init__(self, rank: int):
    """
    Initializes the specification for a low-rank affine flow.

    Parameters:
        rank: The rank R (number of columns in U and V) of the low-rank off-diagonal component.
    """
    self.rank = rank

construct(dim: int) -> LowRankAffineLayer

Constructs a low-rank affine flow layer.

Parameters:

Name Type Description Default
dim int

The dimension of the parameter space.

 required

Returns:

Type Description
LowRankAffineLayer

A LowRankAffineLayer of dimension dim and rank self.rank.
Source code in src/bayinx/flows/lowrankaffine.py 
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def construct(self, dim: int) -> LowRankAffineLayer:
    """
    Constructs a low-rank affine flow layer.

    Parameters:
        dim: The dimension of the parameter space.

    Returns:
        A LowRankAffineLayer of dimension `dim` and rank `self.rank`.
    """
    if (self.rank > (dim - 1)/2):
        raise ValueError(f"Rank {self.rank} is large, consider using a full affine flow instead.")

    return LowRankAffineLayer(dim, self.rank)

Planar

A specification for the Planar flow.

Definition

\(T(\mathbf{z}) = \mathbf{z} + \mathbf{u} h(\mathbf{w}^\top \mathbf{z} + b)\)

Source code in src/bayinx/flows/planar.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
class Planar(FlowSpec):
    """
    A specification for the Planar flow.

    Definition:
        $T(\\mathbf{z}) = \\mathbf{z} + \\mathbf{u} h(\\mathbf{w}^\\top \\mathbf{z} + b)$

    """
    def __init__(self):
        pass

    def construct(self, dim: int) -> PlanarLayer:
        return PlanarLayer(dim)

Sylvester

A specification for the Sylvester flow.

Definition

\(T(\mathbf{z}) = \mathbf{z} + \mathbf{Q} \mathbf{R}_1 h(\mathbf{R}_2 \mathbf{Q}^\top \mathbf{z} + \mathbf{b})\)

Attributes:

Name Type Description
rank int

The rank of the Sylvester flow (number of hidden units).
key PRNGKeyArray

The PRNG key used to generate a Sylvester flow layer.
Source code in src/bayinx/flows/sylvester.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
class Sylvester(FlowSpec):
    """
    A specification for the Sylvester flow.

    Definition:
        $T(\\mathbf{z}) = \\mathbf{z} + \\mathbf{Q} \\mathbf{R}_1 h(\\mathbf{R}_2 \\mathbf{Q}^\\top \\mathbf{z} + \\mathbf{b})$

    Attributes:
        rank: The rank of the Sylvester flow (number of hidden units).
        key: The PRNG key used to generate a Sylvester flow layer.
    """
    rank: int
    key: PRNGKeyArray

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

        Parameters:
            rank: The rank of the transformation (number of hidden units).
            key: A PRNG key used to generate the flow layer.
        """
        self.rank = rank
        self.key = key

    def construct(self, dim: int) -> SylvesterLayer:
        """
        Constructs a Sylvester flow layer.

        Parameters:
            dim: The dimension of the parameter space.

        Returns:
            A SylvesterLayer of dimension `dim` and rank `self.rank`.
        """
        if self.rank > dim:
            raise ValueError(f"Rank {self.rank} cannot be greater than dimension {dim}.")

        return SylvesterLayer(dim, self.rank, self.key)

__init__(rank: int, key: PRNGKeyArray = jr.key(0))

Initializes the specification for a Sylvester flow.

Parameters:

Name Type Description Default
rank int

The rank of the transformation (number of hidden units).

 required
key PRNGKeyArray

A PRNG key used to generate the flow layer.

 key(0)
Source code in src/bayinx/flows/sylvester.py 
222
223
224
225
226
227
228
229
230
231
def __init__(self, rank: int, key: PRNGKeyArray = jr.key(0)):
    """
    Initializes the specification for a Sylvester flow.

    Parameters:
        rank: The rank of the transformation (number of hidden units).
        key: A PRNG key used to generate the flow layer.
    """
    self.rank = rank
    self.key = key

construct(dim: int) -> SylvesterLayer

Constructs a Sylvester flow layer.

Parameters:

Name Type Description Default
dim int

The dimension of the parameter space.

 required

Returns:

Type Description
SylvesterLayer

A SylvesterLayer of dimension dim and rank self.rank.
Source code in src/bayinx/flows/sylvester.py 
233
234
235
236
237
238
239
240
241
242
243
244
245
246
def construct(self, dim: int) -> SylvesterLayer:
    """
    Constructs a Sylvester flow layer.

    Parameters:
        dim: The dimension of the parameter space.

    Returns:
        A SylvesterLayer of dimension `dim` and rank `self.rank`.
    """
    if self.rank > dim:
        raise ValueError(f"Rank {self.rank} cannot be greater than dimension {dim}.")

    return SylvesterLayer(dim, self.rank, self.key)