Skip to content

Distributions

Define the latent distributions used in the normalizing flows.

CentBeta

Bases: LatentDist

A centered Beta distribution.

This distribution is just a regular Beta distribution, scaled and shifted to have support on the domain [-B, B] in each dimension.

Alpha and beta parameters for each dimension are learned during training.

Source code in pzflow/distributions.py 
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 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
127
128
129
130
131
132
133
134
class CentBeta(LatentDist):
    """A centered Beta distribution.

    This distribution is just a regular Beta distribution, scaled and shifted
    to have support on the domain [-B, B] in each dimension.

    Alpha and beta parameters for each dimension are learned during training.
    """

    def __init__(self, input_dim: int, B: float = 5) -> None:
        """
        Parameters
        ----------
        input_dim : int
            The dimension of the distribution.
        B : float; default=5
            The distribution has support (-B, B) along each dimension.
        """
        self.input_dim = input_dim
        self.B = B

        # save dist info
        self._params = tuple([(0.0, 0.0) for i in range(input_dim)])
        self.info = ("CentBeta", (input_dim, B))

    def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
        """Calculates log probability density of inputs.

        Parameters
        ----------
        params : a Jax pytree
            Tuple of ((a1, b1), (a2, b2), ...) where aN,bN are log(alpha),log(beta)
            for the Nth dimension.
        inputs : jnp.ndarray
            Input data for which log probability density is calculated.

        Returns
        -------
        jnp.ndarray
            Device array of shape (inputs.shape[0],).
        """
        log_prob = jnp.hstack(
            [
                beta.logpdf(
                    inputs[:, i],
                    a=jnp.exp(params[i][0]),
                    b=jnp.exp(params[i][1]),
                    loc=-self.B,
                    scale=2 * self.B,
                ).reshape(-1, 1)
                for i in range(self.input_dim)
            ]
        ).sum(axis=1)

        return log_prob

    def sample(
        self, params: Pytree, nsamples: int, seed: int = None
    ) -> jnp.ndarray:
        """Returns samples from the distribution.

        Parameters
        ----------
        params : a Jax pytree
            Tuple of ((a1, b1), (a2, b2), ...) where aN,bN are log(alpha),log(beta)
            for the Nth dimension.
        nsamples : int
            The number of samples to be returned.
        seed : int; optional
            Sets the random seed for the samples.

        Returns
        -------
        jnp.ndarray
            Device array of shape (nsamples, self.input_dim).
        """
        seed = np.random.randint(1e18) if seed is None else seed
        seeds = random.split(random.PRNGKey(seed), self.input_dim)
        samples = jnp.hstack(
            [
                random.beta(
                    seeds[i],
                    jnp.exp(params[i][0]),
                    jnp.exp(params[i][1]),
                    shape=(nsamples, 1),
                )
                for i in range(self.input_dim)
            ]
        )
        return 2 * self.B * (samples - 0.5)

__init__(input_dim, B=5)

Parameters:

Name Type Description Default
input_dim int

The dimension of the distribution.

 required
B float

The distribution has support (-B, B) along each dimension.

 5
Source code in pzflow/distributions.py 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def __init__(self, input_dim: int, B: float = 5) -> None:
    """
    Parameters
    ----------
    input_dim : int
        The dimension of the distribution.
    B : float; default=5
        The distribution has support (-B, B) along each dimension.
    """
    self.input_dim = input_dim
    self.B = B

    # save dist info
    self._params = tuple([(0.0, 0.0) for i in range(input_dim)])
    self.info = ("CentBeta", (input_dim, B))

log_prob(params, inputs)

Calculates log probability density of inputs.

Parameters:

Name Type Description Default
params a Jax pytree

Tuple of ((a1, b1), (a2, b2), ...) where aN,bN are log(alpha),log(beta) for the Nth dimension.

 required
inputs jnp.ndarray

Input data for which log probability density is calculated.

 required

Returns:

Type Description
jnp.ndarray

Device array of shape (inputs.shape[0],).
Source code in pzflow/distributions.py 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
    """Calculates log probability density of inputs.

    Parameters
    ----------
    params : a Jax pytree
        Tuple of ((a1, b1), (a2, b2), ...) where aN,bN are log(alpha),log(beta)
        for the Nth dimension.
    inputs : jnp.ndarray
        Input data for which log probability density is calculated.

    Returns
    -------
    jnp.ndarray
        Device array of shape (inputs.shape[0],).
    """
    log_prob = jnp.hstack(
        [
            beta.logpdf(
                inputs[:, i],
                a=jnp.exp(params[i][0]),
                b=jnp.exp(params[i][1]),
                loc=-self.B,
                scale=2 * self.B,
            ).reshape(-1, 1)
            for i in range(self.input_dim)
        ]
    ).sum(axis=1)

    return log_prob

sample(params, nsamples, seed=None)

Returns samples from the distribution.

Parameters:

Name Type Description Default
params a Jax pytree

Tuple of ((a1, b1), (a2, b2), ...) where aN,bN are log(alpha),log(beta) for the Nth dimension.

 required
nsamples int

The number of samples to be returned.

 required
seed int

Sets the random seed for the samples.

 None

Returns:

Type Description
jnp.ndarray

Device array of shape (nsamples, self.input_dim).
Source code in pzflow/distributions.py 
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
127
128
129
130
131
132
133
134
def sample(
    self, params: Pytree, nsamples: int, seed: int = None
) -> jnp.ndarray:
    """Returns samples from the distribution.

    Parameters
    ----------
    params : a Jax pytree
        Tuple of ((a1, b1), (a2, b2), ...) where aN,bN are log(alpha),log(beta)
        for the Nth dimension.
    nsamples : int
        The number of samples to be returned.
    seed : int; optional
        Sets the random seed for the samples.

    Returns
    -------
    jnp.ndarray
        Device array of shape (nsamples, self.input_dim).
    """
    seed = np.random.randint(1e18) if seed is None else seed
    seeds = random.split(random.PRNGKey(seed), self.input_dim)
    samples = jnp.hstack(
        [
            random.beta(
                seeds[i],
                jnp.exp(params[i][0]),
                jnp.exp(params[i][1]),
                shape=(nsamples, 1),
            )
            for i in range(self.input_dim)
        ]
    )
    return 2 * self.B * (samples - 0.5)

CentBeta13

Bases: LatentDist

A centered Beta distribution with alpha, beta = 13.

This distribution is just a regular Beta distribution, scaled and shifted to have support on the domain [-B, B] in each dimension.

Alpha, beta = 13 means that the distribution looks like a Gaussian distribution, but with hard cutoffs at +/- B.

Source code in pzflow/distributions.py 
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
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
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
class CentBeta13(LatentDist):
    """A centered Beta distribution with alpha, beta = 13.

    This distribution is just a regular Beta distribution, scaled and shifted
    to have support on the domain [-B, B] in each dimension.

    Alpha, beta = 13 means that the distribution looks like a Gaussian
    distribution, but with hard cutoffs at +/- B.
    """

    def __init__(self, input_dim: int, B: float = 5) -> None:
        """
        Parameters
        ----------
        input_dim : int
            The dimension of the distribution.
        B : float; default=5
            The distribution has support (-B, B) along each dimension.
        """
        self.input_dim = input_dim
        self.B = B

        # save dist info
        self._params = tuple([(0.0, 0.0) for i in range(input_dim)])
        self.info = ("CentBeta13", (input_dim, B))
        self.a = 13
        self.b = 13

    def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
        """Calculates log probability density of inputs.

        Parameters
        ----------
        params : a Jax pytree
            Empty pytree -- this distribution doesn't have learnable parameters.
            This parameter is present to ensure a consistent interface.
        inputs : jnp.ndarray
            Input data for which log probability density is calculated.

        Returns
        -------
        jnp.ndarray
            Device array of shape (inputs.shape[0],).
        """
        log_prob = jnp.hstack(
            [
                beta.logpdf(
                    inputs[:, i],
                    a=self.a,
                    b=self.b,
                    loc=-self.B,
                    scale=2 * self.B,
                ).reshape(-1, 1)
                for i in range(self.input_dim)
            ]
        ).sum(axis=1)

        return log_prob

    def sample(
        self, params: Pytree, nsamples: int, seed: int = None
    ) -> jnp.ndarray:
        """Returns samples from the distribution.

        Parameters
        ----------
        params : a Jax pytree
            Empty pytree -- this distribution doesn't have learnable parameters.
            This parameter is present to ensure a consistent interface.
        nsamples : int
            The number of samples to be returned.
        seed : int; optional
            Sets the random seed for the samples.

        Returns
        -------
        jnp.ndarray
            Device array of shape (nsamples, self.input_dim).
        """
        seed = np.random.randint(1e18) if seed is None else seed
        seeds = random.split(random.PRNGKey(seed), self.input_dim)
        samples = jnp.hstack(
            [
                random.beta(
                    seeds[i],
                    self.a,
                    self.b,
                    shape=(nsamples, 1),
                )
                for i in range(self.input_dim)
            ]
        )
        return 2 * self.B * (samples - 0.5)

__init__(input_dim, B=5)

Parameters:

Name Type Description Default
input_dim int

The dimension of the distribution.

 required
B float

The distribution has support (-B, B) along each dimension.

 5
Source code in pzflow/distributions.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def __init__(self, input_dim: int, B: float = 5) -> None:
    """
    Parameters
    ----------
    input_dim : int
        The dimension of the distribution.
    B : float; default=5
        The distribution has support (-B, B) along each dimension.
    """
    self.input_dim = input_dim
    self.B = B

    # save dist info
    self._params = tuple([(0.0, 0.0) for i in range(input_dim)])
    self.info = ("CentBeta13", (input_dim, B))
    self.a = 13
    self.b = 13

log_prob(params, inputs)

Calculates log probability density of inputs.

Parameters:

Name Type Description Default
params a Jax pytree

Empty pytree -- this distribution doesn't have learnable parameters. This parameter is present to ensure a consistent interface.

 required
inputs jnp.ndarray

Input data for which log probability density is calculated.

 required

Returns:

Type Description
jnp.ndarray

Device array of shape (inputs.shape[0],).
Source code in pzflow/distributions.py 
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
def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
    """Calculates log probability density of inputs.

    Parameters
    ----------
    params : a Jax pytree
        Empty pytree -- this distribution doesn't have learnable parameters.
        This parameter is present to ensure a consistent interface.
    inputs : jnp.ndarray
        Input data for which log probability density is calculated.

    Returns
    -------
    jnp.ndarray
        Device array of shape (inputs.shape[0],).
    """
    log_prob = jnp.hstack(
        [
            beta.logpdf(
                inputs[:, i],
                a=self.a,
                b=self.b,
                loc=-self.B,
                scale=2 * self.B,
            ).reshape(-1, 1)
            for i in range(self.input_dim)
        ]
    ).sum(axis=1)

    return log_prob

sample(params, nsamples, seed=None)

Returns samples from the distribution.

Parameters:

Name Type Description Default
params a Jax pytree

Empty pytree -- this distribution doesn't have learnable parameters. This parameter is present to ensure a consistent interface.

 required
nsamples int

The number of samples to be returned.

 required
seed int

Sets the random seed for the samples.

 None

Returns:

Type Description
jnp.ndarray

Device array of shape (nsamples, self.input_dim).
Source code in pzflow/distributions.py 
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
def sample(
    self, params: Pytree, nsamples: int, seed: int = None
) -> jnp.ndarray:
    """Returns samples from the distribution.

    Parameters
    ----------
    params : a Jax pytree
        Empty pytree -- this distribution doesn't have learnable parameters.
        This parameter is present to ensure a consistent interface.
    nsamples : int
        The number of samples to be returned.
    seed : int; optional
        Sets the random seed for the samples.

    Returns
    -------
    jnp.ndarray
        Device array of shape (nsamples, self.input_dim).
    """
    seed = np.random.randint(1e18) if seed is None else seed
    seeds = random.split(random.PRNGKey(seed), self.input_dim)
    samples = jnp.hstack(
        [
            random.beta(
                seeds[i],
                self.a,
                self.b,
                shape=(nsamples, 1),
            )
            for i in range(self.input_dim)
        ]
    )
    return 2 * self.B * (samples - 0.5)

Joint

Bases: LatentDist

A joint distribution built from other distributions.

Note that each of the other distributions already have support for multiple dimensions. This is only useful if you want to combine different distributions for different dimensions, e.g. if your first dimension has a Uniform latent space and the second dimension has a CentBeta latent space.

Source code in pzflow/distributions.py 
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
class Joint(LatentDist):
    """A joint distribution built from other distributions.

    Note that each of the other distributions already have support for
    multiple dimensions. This is only useful if you want to combine
    different distributions for different dimensions, e.g. if your first
    dimension has a Uniform latent space and the second dimension has a
    CentBeta latent space.
    """

    def __init__(self, *inputs: Union[LatentDist, tuple]) -> None:
        """
        Parameters
        ----------
        inputs: LatentDist or tuple
            The latent distributions to join together.
        """

        # if Joint info provided, use that for setup
        if inputs[0] == "Joint info":
            self.dists = [globals()[dist[0]](*dist[1]) for dist in inputs[1]]
        # otherwise, assume it's a list of distributions
        else:
            self.dists = inputs

        # save info
        self._params = [dist._params for dist in self.dists]
        self.input_dim = sum([dist.input_dim for dist in self.dists])
        self.info = (
            "Joint",
            ("Joint info", [dist.info for dist in self.dists]),
        )

        # save the indices at which inputs will be split for log_prob
        # they must be concretely saved ahead-of-time so that jax trace
        # works properly when jitting
        self._splits = jnp.cumsum(
            jnp.array([dist.input_dim for dist in self.dists])
        )[:-1]

    def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
        """Calculates log probability density of inputs.

        Parameters
        ----------
        params : Jax Pytree
            Parameters for the distributions.
        inputs : jnp.ndarray
            Input data for which log probability density is calculated.

        Returns
        -------
        jnp.ndarray
            Device array of shape (inputs.shape[0],).
        """

        # split inputs for corresponding distribution
        inputs = jnp.split(inputs, self._splits, axis=1)

        # calculate log_prob with respect to each sub-distribution,
        # then sum all the log_probs for each input
        log_prob = jnp.hstack(
            [
                self.dists[i].log_prob(params[i], inputs[i]).reshape(-1, 1)
                for i in range(len(self.dists))
            ]
        ).sum(axis=1)

        return log_prob

    def sample(
        self, params: Pytree, nsamples: int, seed: int = None
    ) -> jnp.ndarray:
        """Returns samples from the distribution.

        Parameters
        ----------
        params : a Jax pytree
            Parameters for the distributions.
        nsamples : int
            The number of samples to be returned.
        seed : int; optional
            Sets the random seed for the samples.

        Returns
        -------
        jnp.ndarray
            Device array of shape (nsamples, self.input_dim).
        """

        seed = np.random.randint(1e18) if seed is None else seed
        seeds = random.randint(
            random.PRNGKey(seed), (len(self.dists),), 0, int(1e9)
        )
        samples = jnp.hstack(
            [
                self.dists[i]
                .sample(params[i], nsamples, seeds[i])
                .reshape(nsamples, -1)
                for i in range(len(self.dists))
            ]
        )

        return samples

__init__(inputs)

Parameters:

Name Type Description Default
inputs Union[LatentDist, tuple]

The latent distributions to join together.

 ()
Source code in pzflow/distributions.py 
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
def __init__(self, *inputs: Union[LatentDist, tuple]) -> None:
    """
    Parameters
    ----------
    inputs: LatentDist or tuple
        The latent distributions to join together.
    """

    # if Joint info provided, use that for setup
    if inputs[0] == "Joint info":
        self.dists = [globals()[dist[0]](*dist[1]) for dist in inputs[1]]
    # otherwise, assume it's a list of distributions
    else:
        self.dists = inputs

    # save info
    self._params = [dist._params for dist in self.dists]
    self.input_dim = sum([dist.input_dim for dist in self.dists])
    self.info = (
        "Joint",
        ("Joint info", [dist.info for dist in self.dists]),
    )

    # save the indices at which inputs will be split for log_prob
    # they must be concretely saved ahead-of-time so that jax trace
    # works properly when jitting
    self._splits = jnp.cumsum(
        jnp.array([dist.input_dim for dist in self.dists])
    )[:-1]

log_prob(params, inputs)

Calculates log probability density of inputs.

Parameters:

Name Type Description Default
params Jax Pytree

Parameters for the distributions.

 required
inputs jnp.ndarray

Input data for which log probability density is calculated.

 required

Returns:

Type Description
jnp.ndarray

Device array of shape (inputs.shape[0],).
Source code in pzflow/distributions.py 
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
    """Calculates log probability density of inputs.

    Parameters
    ----------
    params : Jax Pytree
        Parameters for the distributions.
    inputs : jnp.ndarray
        Input data for which log probability density is calculated.

    Returns
    -------
    jnp.ndarray
        Device array of shape (inputs.shape[0],).
    """

    # split inputs for corresponding distribution
    inputs = jnp.split(inputs, self._splits, axis=1)

    # calculate log_prob with respect to each sub-distribution,
    # then sum all the log_probs for each input
    log_prob = jnp.hstack(
        [
            self.dists[i].log_prob(params[i], inputs[i]).reshape(-1, 1)
            for i in range(len(self.dists))
        ]
    ).sum(axis=1)

    return log_prob

sample(params, nsamples, seed=None)

Returns samples from the distribution.

Parameters:

Name Type Description Default
params a Jax pytree

Parameters for the distributions.

 required
nsamples int

The number of samples to be returned.

 required
seed int

Sets the random seed for the samples.

 None

Returns:

Type Description
jnp.ndarray

Device array of shape (nsamples, self.input_dim).
Source code in pzflow/distributions.py 
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
def sample(
    self, params: Pytree, nsamples: int, seed: int = None
) -> jnp.ndarray:
    """Returns samples from the distribution.

    Parameters
    ----------
    params : a Jax pytree
        Parameters for the distributions.
    nsamples : int
        The number of samples to be returned.
    seed : int; optional
        Sets the random seed for the samples.

    Returns
    -------
    jnp.ndarray
        Device array of shape (nsamples, self.input_dim).
    """

    seed = np.random.randint(1e18) if seed is None else seed
    seeds = random.randint(
        random.PRNGKey(seed), (len(self.dists),), 0, int(1e9)
    )
    samples = jnp.hstack(
        [
            self.dists[i]
            .sample(params[i], nsamples, seeds[i])
            .reshape(nsamples, -1)
            for i in range(len(self.dists))
        ]
    )

    return samples

LatentDist

Bases: ABC

Base class for latent distributions.

Source code in pzflow/distributions.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class LatentDist(ABC):
    """Base class for latent distributions."""

    info = ("LatentDist", ())

    @abstractmethod
    def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
        """Calculate log-probability of the inputs."""

    @abstractmethod
    def sample(
        self, params: Pytree, nsamples: int, seed: int = None
    ) -> jnp.ndarray:
        """Sample from the distribution."""

log_prob(params, inputs) abstractmethod

Calculate log-probability of the inputs.

Source code in pzflow/distributions.py 
22
23
24
@abstractmethod
def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
    """Calculate log-probability of the inputs."""

sample(params, nsamples, seed=None) abstractmethod

Sample from the distribution.

Source code in pzflow/distributions.py 
26
27
28
29
30
@abstractmethod
def sample(
    self, params: Pytree, nsamples: int, seed: int = None
) -> jnp.ndarray:
    """Sample from the distribution."""

Normal

Bases: LatentDist

A multivariate Gaussian distribution with mean zero and unit variance.

Note this distribution has infinite support, so it is not recommended that you use it with the spline coupling layers, which have compact support. If you do use the two together, you should set the support of the spline layers (using the spline parameter B) to be large enough that you rarely draw Gaussian samples outside the support of the splines.

Source code in pzflow/distributions.py 
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
class Normal(LatentDist):
    """A multivariate Gaussian distribution with mean zero and unit variance.

    Note this distribution has infinite support, so it is not recommended that
    you use it with the spline coupling layers, which have compact support.
    If you do use the two together, you should set the support of the spline
    layers (using the spline parameter B) to be large enough that you rarely
    draw Gaussian samples outside the support of the splines.
    """

    def __init__(self, input_dim: int) -> None:
        """
        Parameters
        ----------
        input_dim : int
            The dimension of the distribution.
        """
        self.input_dim = input_dim

        # save dist info
        self._params = ()
        self.info = ("Normal", (input_dim,))

    def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
        """Calculates log probability density of inputs.

        Parameters
        ----------
        params : a Jax pytree
            Empty pytree -- this distribution doesn't have learnable parameters.
            This parameter is present to ensure a consistent interface.
        inputs : jnp.ndarray
            Input data for which log probability density is calculated.

        Returns
        -------
        jnp.ndarray
            Device array of shape (inputs.shape[0],).
        """
        return multivariate_normal.logpdf(
            inputs,
            mean=jnp.zeros(self.input_dim),
            cov=jnp.identity(self.input_dim),
        )

    def sample(
        self, params: Pytree, nsamples: int, seed: int = None
    ) -> jnp.ndarray:
        """Returns samples from the distribution.

        Parameters
        ----------
        params : a Jax pytree
            Empty pytree -- this distribution doesn't have learnable parameters.
            This parameter is present to ensure a consistent interface.
        nsamples : int
            The number of samples to be returned.
        seed : int; optional
            Sets the random seed for the samples.

        Returns
        -------
        jnp.ndarray
            Device array of shape (nsamples, self.input_dim).
        """
        seed = np.random.randint(1e18) if seed is None else seed
        return random.multivariate_normal(
            key=random.PRNGKey(seed),
            mean=jnp.zeros(self.input_dim),
            cov=jnp.identity(self.input_dim),
            shape=(nsamples,),
        )

__init__(input_dim)

Parameters:

Name Type Description Default
input_dim int

The dimension of the distribution.

 required
Source code in pzflow/distributions.py 
242
243
244
245
246
247
248
249
250
251
252
253
def __init__(self, input_dim: int) -> None:
    """
    Parameters
    ----------
    input_dim : int
        The dimension of the distribution.
    """
    self.input_dim = input_dim

    # save dist info
    self._params = ()
    self.info = ("Normal", (input_dim,))

log_prob(params, inputs)

Calculates log probability density of inputs.

Parameters:

Name Type Description Default
params a Jax pytree

Empty pytree -- this distribution doesn't have learnable parameters. This parameter is present to ensure a consistent interface.

 required
inputs jnp.ndarray

Input data for which log probability density is calculated.

 required

Returns:

Type Description
jnp.ndarray

Device array of shape (inputs.shape[0],).
Source code in pzflow/distributions.py 
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
    """Calculates log probability density of inputs.

    Parameters
    ----------
    params : a Jax pytree
        Empty pytree -- this distribution doesn't have learnable parameters.
        This parameter is present to ensure a consistent interface.
    inputs : jnp.ndarray
        Input data for which log probability density is calculated.

    Returns
    -------
    jnp.ndarray
        Device array of shape (inputs.shape[0],).
    """
    return multivariate_normal.logpdf(
        inputs,
        mean=jnp.zeros(self.input_dim),
        cov=jnp.identity(self.input_dim),
    )

sample(params, nsamples, seed=None)

Returns samples from the distribution.

Parameters:

Name Type Description Default
params a Jax pytree

Empty pytree -- this distribution doesn't have learnable parameters. This parameter is present to ensure a consistent interface.

 required
nsamples int

The number of samples to be returned.

 required
seed int

Sets the random seed for the samples.

 None

Returns:

Type Description
jnp.ndarray

Device array of shape (nsamples, self.input_dim).
Source code in pzflow/distributions.py 
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
def sample(
    self, params: Pytree, nsamples: int, seed: int = None
) -> jnp.ndarray:
    """Returns samples from the distribution.

    Parameters
    ----------
    params : a Jax pytree
        Empty pytree -- this distribution doesn't have learnable parameters.
        This parameter is present to ensure a consistent interface.
    nsamples : int
        The number of samples to be returned.
    seed : int; optional
        Sets the random seed for the samples.

    Returns
    -------
    jnp.ndarray
        Device array of shape (nsamples, self.input_dim).
    """
    seed = np.random.randint(1e18) if seed is None else seed
    return random.multivariate_normal(
        key=random.PRNGKey(seed),
        mean=jnp.zeros(self.input_dim),
        cov=jnp.identity(self.input_dim),
        shape=(nsamples,),
    )

Tdist

Bases: LatentDist

A multivariate T distribution with mean zero and unit scale matrix.

The number of degrees of freedom (i.e. the weight of the tails) is learned during training.

Note this distribution has infinite support and potentially large tails, so it is not recommended to use this distribution with the spline coupling layers, which have compact support.

Source code in pzflow/distributions.py 
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
class Tdist(LatentDist):
    """A multivariate T distribution with mean zero and unit scale matrix.

    The number of degrees of freedom (i.e. the weight of the tails) is learned
    during training.

    Note this distribution has infinite support and potentially large tails,
    so it is not recommended to use this distribution with the spline coupling
    layers, which have compact support.
    """

    def __init__(self, input_dim: int) -> None:
        """
        Parameters
        ----------
        input_dim : int
            The dimension of the distribution.
        """
        self.input_dim = input_dim

        # save dist info
        self._params = jnp.log(30.0)
        self.info = ("Tdist", (input_dim,))

    def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
        """Calculates log probability density of inputs.

        Uses method explained here:
        http://gregorygundersen.com/blog/2020/01/20/multivariate-t/

        Parameters
        ----------
        params : float
            The degrees of freedom (nu) of the t-distribution.
        inputs : jnp.ndarray
            Input data for which log probability density is calculated.

        Returns
        -------
        jnp.ndarray
            Device array of shape (inputs.shape[0],).
        """
        cov = jnp.identity(self.input_dim)
        nu = jnp.exp(params)
        maha, log_det = _mahalanobis_and_logdet(inputs, cov)
        t = 0.5 * (nu + self.input_dim)
        A = gammaln(t)
        B = gammaln(0.5 * nu)
        C = self.input_dim / 2.0 * jnp.log(nu * jnp.pi)
        D = 0.5 * log_det
        E = -t * jnp.log(1 + (1.0 / nu) * maha)

        return A - B - C - D + E

    def sample(
        self, params: Pytree, nsamples: int, seed: int = None
    ) -> jnp.ndarray:
        """Returns samples from the distribution.

        Parameters
        ----------
        params : float
            The degrees of freedom (nu) of the t-distribution.
        nsamples : int
            The number of samples to be returned.
        seed : int; optional
            Sets the random seed for the samples.

        Returns
        -------
        jnp.ndarray
            Device array of shape (nsamples, self.input_dim).
        """
        mean = jnp.zeros(self.input_dim)
        nu = jnp.exp(params)

        seed = np.random.randint(1e18) if seed is None else seed
        rng = np.random.default_rng(int(seed))
        x = jnp.array(rng.chisquare(nu, nsamples) / nu)
        z = random.multivariate_normal(
            key=random.PRNGKey(seed),
            mean=jnp.zeros(self.input_dim),
            cov=jnp.identity(self.input_dim),
            shape=(nsamples,),
        )
        samples = mean + z / jnp.sqrt(x)[:, None]
        return samples

__init__(input_dim)

Parameters:

Name Type Description Default
input_dim int

The dimension of the distribution.

 required
Source code in pzflow/distributions.py 
317
318
319
320
321
322
323
324
325
326
327
328
def __init__(self, input_dim: int) -> None:
    """
    Parameters
    ----------
    input_dim : int
        The dimension of the distribution.
    """
    self.input_dim = input_dim

    # save dist info
    self._params = jnp.log(30.0)
    self.info = ("Tdist", (input_dim,))

log_prob(params, inputs)

Calculates log probability density of inputs.

Uses method explained here: http://gregorygundersen.com/blog/2020/01/20/multivariate-t/

Parameters:

Name Type Description Default
params float

The degrees of freedom (nu) of the t-distribution.

 required
inputs jnp.ndarray

Input data for which log probability density is calculated.

 required

Returns:

Type Description
jnp.ndarray

Device array of shape (inputs.shape[0],).
Source code in pzflow/distributions.py 
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
    """Calculates log probability density of inputs.

    Uses method explained here:
    http://gregorygundersen.com/blog/2020/01/20/multivariate-t/

    Parameters
    ----------
    params : float
        The degrees of freedom (nu) of the t-distribution.
    inputs : jnp.ndarray
        Input data for which log probability density is calculated.

    Returns
    -------
    jnp.ndarray
        Device array of shape (inputs.shape[0],).
    """
    cov = jnp.identity(self.input_dim)
    nu = jnp.exp(params)
    maha, log_det = _mahalanobis_and_logdet(inputs, cov)
    t = 0.5 * (nu + self.input_dim)
    A = gammaln(t)
    B = gammaln(0.5 * nu)
    C = self.input_dim / 2.0 * jnp.log(nu * jnp.pi)
    D = 0.5 * log_det
    E = -t * jnp.log(1 + (1.0 / nu) * maha)

    return A - B - C - D + E

sample(params, nsamples, seed=None)

Returns samples from the distribution.

Parameters:

Name Type Description Default
params float

The degrees of freedom (nu) of the t-distribution.

 required
nsamples int

The number of samples to be returned.

 required
seed int

Sets the random seed for the samples.

 None

Returns:

Type Description
jnp.ndarray

Device array of shape (nsamples, self.input_dim).
Source code in pzflow/distributions.py 
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
def sample(
    self, params: Pytree, nsamples: int, seed: int = None
) -> jnp.ndarray:
    """Returns samples from the distribution.

    Parameters
    ----------
    params : float
        The degrees of freedom (nu) of the t-distribution.
    nsamples : int
        The number of samples to be returned.
    seed : int; optional
        Sets the random seed for the samples.

    Returns
    -------
    jnp.ndarray
        Device array of shape (nsamples, self.input_dim).
    """
    mean = jnp.zeros(self.input_dim)
    nu = jnp.exp(params)

    seed = np.random.randint(1e18) if seed is None else seed
    rng = np.random.default_rng(int(seed))
    x = jnp.array(rng.chisquare(nu, nsamples) / nu)
    z = random.multivariate_normal(
        key=random.PRNGKey(seed),
        mean=jnp.zeros(self.input_dim),
        cov=jnp.identity(self.input_dim),
        shape=(nsamples,),
    )
    samples = mean + z / jnp.sqrt(x)[:, None]
    return samples

Uniform

Bases: LatentDist

A multivariate uniform distribution with support [-B, B].

Source code in pzflow/distributions.py 
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
class Uniform(LatentDist):
    """A multivariate uniform distribution with support [-B, B]."""

    def __init__(self, input_dim: int, B: float = 5) -> None:
        """
        Parameters
        ----------
        input_dim : int
            The dimension of the distribution.
        B : float; default=5
            The distribution has support (-B, B) along each dimension.
        """
        self.input_dim = input_dim
        self.B = B

        # save dist info
        self._params = ()
        self.info = ("Uniform", (input_dim, B))

    def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
        """Calculates log probability density of inputs.

        Parameters
        ----------
        params : Jax Pytree
            Empty pytree -- this distribution doesn't have learnable parameters.
            This parameter is present to ensure a consistent interface.
        inputs : jnp.ndarray
            Input data for which log probability density is calculated.

        Returns
        -------
        jnp.ndarray
            Device array of shape (inputs.shape[0],).
        """

        # which inputs are inside the support of the distribution
        mask = jnp.prod((inputs >= -self.B) & (inputs <= self.B), axis=-1)

        # calculate log_prob
        log_prob = jnp.where(
            mask,
            -self.input_dim * jnp.log(2 * self.B),
            -jnp.inf,
        )

        return log_prob

    def sample(
        self, params: Pytree, nsamples: int, seed: int = None
    ) -> jnp.ndarray:
        """Returns samples from the distribution.

        Parameters
        ----------
        params : a Jax pytree
            Empty pytree -- this distribution doesn't have learnable parameters.
            This parameter is present to ensure a consistent interface.
        nsamples : int
            The number of samples to be returned.
        seed : int; optional
            Sets the random seed for the samples.

        Returns
        -------
        jnp.ndarray
            Device array of shape (nsamples, self.input_dim).
        """
        seed = np.random.randint(1e18) if seed is None else seed
        samples = random.uniform(
            random.PRNGKey(seed),
            shape=(nsamples, self.input_dim),
            minval=-self.B,
            maxval=self.B,
        )
        return jnp.array(samples)

__init__(input_dim, B=5)

Parameters:

Name Type Description Default
input_dim int

The dimension of the distribution.

 required
B float

The distribution has support (-B, B) along each dimension.

 5
Source code in pzflow/distributions.py 
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
def __init__(self, input_dim: int, B: float = 5) -> None:
    """
    Parameters
    ----------
    input_dim : int
        The dimension of the distribution.
    B : float; default=5
        The distribution has support (-B, B) along each dimension.
    """
    self.input_dim = input_dim
    self.B = B

    # save dist info
    self._params = ()
    self.info = ("Uniform", (input_dim, B))

log_prob(params, inputs)

Calculates log probability density of inputs.

Parameters:

Name Type Description Default
params Jax Pytree

Empty pytree -- this distribution doesn't have learnable parameters. This parameter is present to ensure a consistent interface.

 required
inputs jnp.ndarray

Input data for which log probability density is calculated.

 required

Returns:

Type Description
jnp.ndarray

Device array of shape (inputs.shape[0],).
Source code in pzflow/distributions.py 
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
def log_prob(self, params: Pytree, inputs: jnp.ndarray) -> jnp.ndarray:
    """Calculates log probability density of inputs.

    Parameters
    ----------
    params : Jax Pytree
        Empty pytree -- this distribution doesn't have learnable parameters.
        This parameter is present to ensure a consistent interface.
    inputs : jnp.ndarray
        Input data for which log probability density is calculated.

    Returns
    -------
    jnp.ndarray
        Device array of shape (inputs.shape[0],).
    """

    # which inputs are inside the support of the distribution
    mask = jnp.prod((inputs >= -self.B) & (inputs <= self.B), axis=-1)

    # calculate log_prob
    log_prob = jnp.where(
        mask,
        -self.input_dim * jnp.log(2 * self.B),
        -jnp.inf,
    )

    return log_prob

sample(params, nsamples, seed=None)

Returns samples from the distribution.

Parameters:

Name Type Description Default
params a Jax pytree

Empty pytree -- this distribution doesn't have learnable parameters. This parameter is present to ensure a consistent interface.

 required
nsamples int

The number of samples to be returned.

 required
seed int

Sets the random seed for the samples.

 None

Returns:

Type Description
jnp.ndarray

Device array of shape (nsamples, self.input_dim).
Source code in pzflow/distributions.py 
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
def sample(
    self, params: Pytree, nsamples: int, seed: int = None
) -> jnp.ndarray:
    """Returns samples from the distribution.

    Parameters
    ----------
    params : a Jax pytree
        Empty pytree -- this distribution doesn't have learnable parameters.
        This parameter is present to ensure a consistent interface.
    nsamples : int
        The number of samples to be returned.
    seed : int; optional
        Sets the random seed for the samples.

    Returns
    -------
    jnp.ndarray
        Device array of shape (nsamples, self.input_dim).
    """
    seed = np.random.randint(1e18) if seed is None else seed
    samples = random.uniform(
        random.PRNGKey(seed),
        shape=(nsamples, self.input_dim),
        minval=-self.B,
        maxval=self.B,
    )
    return jnp.array(samples)