U1.L1 — Coupling Flows on 2D Toys

Flow-Based Generative Models · UFRJ · 2026.2

01Orientation and handover

Where this session sits

Session	What happened
U1.T1	a coupling layer has a cheap inverse and a log-det that is a sum
U1.L1 — today	you write it, check it, train it, look at it, read the library’s version
U1.T2	what such a stack can and cannot represent

Today is the first lab of the joint phase, and the first time a generative model is written from nothing and trained.

Zero to ten minutes for block A. Say the joint-audience sentence once: this lab depends on U1.T1 and on nothing else. Three modules arrive as given code; their contracts are on the next slide, and the bootcamp is provenance, not a prerequisite. Notebook open on the projector from here on; the deck is the map, the notebook is the road.

Three objects you leave with

1. A bijection you wrote — AffineCoupling, and a Flow that stacks it.
2. Exact density evaluation — Flow.log_prob, the change-of-variables theorem as code.
3. An independent Jacobian witness — autodiff, per example, float64, before and after training.

And one thing deliberately not on the list: an explanation of the last picture of the day.

Write the three objects on the board, top-left, and tick them off as the blocks pass. The fourth item is the U1.T2 hook; do not preview it.

Given code, three modules

Module	Built in	What this lab uses
common/train.py	U0.L1, U0.L2	set_seed; Trainer and its .fit
common/toy2d.py	today	three frozen generators, two picture helpers
common/coupling.py	today	the reference you rebuild

The Trainer: five-step loop, schedule stepped per step, validation under no_grad every epoch, last.pt and best.pt, a CSV log. It expects (x, y) batches; a flow has no labels, so y is a placeholder and the accuracy reads nan — an interface preserved, not a defect.

The parity rule, executable: a joint session may CONSUME bootcamp artifacts as given code with a contract-level description, and may never REQUIRE having attended. This table is the pattern every later joint lab copies. PGMAT readers: this is your loop from U0.L1 with the loss line swapped.

Where the curves go

MLflow, deployed locally: sqlite:///mlflow.db next to the notebook, plus mlartifacts/. Portable — an offline bundle stands alone. Not a production recommendation.

mlflow ui --backend-store-uri sqlite:///mlflow.db

One departure, measured: per-step logging turned a 14 s training into 48 s (a step is under a millisecond; a SQLite write is a few). So the loop runs without the tracker, and the two NLL curves (negative log-likelihood: the loss, on train and on validation) are logged per epoch, after the loop, with the config and the tags.

The U0.L2 decision, stated accurately: SQLite because the mlruns/ file store is deprecated in MLflow 3.15; local because a teaching bundle must be portable. The per-epoch departure is the honest engineering note: name the numbers.

The sanity ritual, third appearance

\[ v^{\top}\nabla_\theta L \;\approx\; \frac{L(\theta+\varepsilon v)-L(\theta-\varepsilon v)}{2\varepsilon}, \qquad \lVert v \rVert = 1 \]

Float64, one fixed unit direction, mixed tolerance \(\texttt{atol} + \texttt{rtol}\cdot\max(|d_{\mathrm{AD}}|, |d_{\mathrm{FD}}|)\). Two extra forward evaluations, whatever the parameter count.

Today’s sibling: the suspect is a log-determinant. Analytic sum of scales, against a Jacobian that autodiff assembles independently, per example.

Run the scalar-toy cell live (three lines). Then say the sibling sentence and move on; block C is where it is executed.

02toy2d: the datasets, frozen today

Three generators, one contract

Ten to eighteen minutes. Infrastructure, not a TODO. Run the plotting cell, show the three panels, then the contract slide. Point at the middle panel: eight squares, outlined, nothing between them.

The contract

two_moons / checkerboard / rings (n, seed)
    -> Tensor (n, 2), float32;  EXTENT = (-4.0, 4.0)
density_heatmap(log_prob_fn, extent, n, ax)
    -> (xs, ys, logp)
warped_grid(generative_fn, extent, ax) -> lines

1. Pure seeded functions. Same seed, same pinned environment, same tensor, bit for bit.
2. Float32, shape (n, 2).
3. Literal zero-density gaps in the checkerboard: eight separated squares of side \(1.5\), no jitter.

Downstream sessions import this module by name and never redefine a generator.

Why frozen: U2.L2 and U3.L1 train other models on exactly these points and compare against today’s. A comparison on different data is not a comparison. Checkpoint 1 runs the contract: bit-identical, seed-sensitive, 100,000 checkerboard points with 0 outside the squares.

The gapped checkerboard is a witness

• Mixture of eight squares of side \(1.5\), centred in the black cells of a \(4\times4\) board on \([-4,4]^2\).
• Nearest corners of two same-colour squares: \(0.5\sqrt2 \approx 0.71\) apart. White cells: no mass at all.
• The support of \(q\) is a union of eight disjoint closed squares; the density is exactly \(1/18\) on it.

That is a mathematical statement about literal support, and it is there because U1.T2 needs one. Today you only observe what a fitted density does in the gaps.

Say “literal support” twice. Do not say why it matters beyond “U1.T2 needs a witness”. The number 1/18 returns in block D as the NLL the model cannot reach.

03An affine coupling layer, from scratch

The layer, generative direction

\[ \begin{aligned} \text{forward:}\; & y_A = x_A,\quad y_B = x_B \odot e^{s_\theta(x_A)} + t_\theta(x_A), \\ & \log|\det J_\psi| = \textstyle\sum_j s_\theta(x_A)_j \\[4pt] \text{inverse:}\; & x_A = y_A,\quad x_B = (y_B - t_\theta(y_A)) \odot e^{-s_\theta(y_A)}, \\ & \log|\det J_{\psi^{-1}}| = -\textstyle\sum_j s_\theta(y_A)_j \end{aligned} \]

Both sums come off the same conditioner pass: \(x_A = y_A\) is on both sides. In code, z is the source-side input (\(z_{k-1}\) of the T1 chain; \(z_0 = x_0\)), x the data-side output, mask marks \(x_A\) with True.

Eighteen to thirty-eight minutes for block C. BOARD 2, left half: the two equations, the sign of the inverse log-det circled. Names: z is source-side, x is data-side, per layer; do not let anyone call a coordinate x_1.

The interface is frozen

AffineCoupling(d, mask, hidden=(64, 64), s_cap=2.0)
    forward(z) -> (x, logdet_fwd)  # psi: source -> data
    inverse(x) -> (z, logdet_inv)  # psi^-1: data -> src
    # logdet_inv == -logdet_fwd, PER EXAMPLE
Flow(layers): forward / inverse / log_prob(x) / sample(n)
Permute(perm): both log-dets exactly 0

TODO 2.1 writes _heads, forward, inverse; the constructor and the conditioner (an MLP \(\mathbb{R}^{d_A} \to \mathbb{R}^{2 d_B}\), two heads) are given. Checkpoint: x[:, mask] == z[:, mask], and the log-det at init is not zero.

Live-code forward once, slowly: split by the mask, heads, scale-and-shift, sum s. Then let them write inverse. Per-example log-dets in BOTH directions is the contract; a library that returns one gets translated to this, never the reverse.

Bound the scale — the policy

\[ s \;=\; s_{\max}\tanh\!\big(\text{raw}/s_{\max}\big), \qquad s_{\max} = 2 \]

So \(|s| < 2\): every layer’s volume factor \(e^{s}\) stays in \([0.14,\, 7.4]\). Near zero the map is the identity; far from zero it saturates, so a random conditioner can neither overflow \(e^{s}\) nor divide the inverse by \(\approx 0\).

U1.T1 promised to show what the loss looks like when you forget. Next slide, measured.

Zuko does the same with a soft clip at |s| < 6.9 — block F reads it. Same policy, different shape and constant.

Bound the scale — the measurement

Three seeds each, 100 epochs. Left, the course default learning rate: six curves on top of each other, the bound changes nothing you can see. Right, learning rate 1e-2: two of three unbounded runs produce NaN and stop (the crosses); one bounded run spikes off the axis near epoch 30 and recovers; no bounded run died in the twelve runs. Insurance is invisible until it pays. One line, one tanh.

Why the head is not zero

• The recipe card zero-inits a last layer where “output \(\approx 0\) at start” is sensible — and an identity-at-start coupling layer is sensible.
• But a zero head gives \(s \equiv 0\): analytic log-det \(0\), autodiff log-det \(0\), perfect agreement on nothing. U0.L3 and U0.L4 both met this.
• So the head starts small and nonzero (default init \(\times\, 0.1\)). Measured at init, six layers, 64 points: \(\max|\log|\det J|| = 0.020\).

A check that cannot fail has verified nothing. The notebook asserts the log-det is nonzero before trusting any agreement.

This is the third sighting of the same species in the course; name it as such. Zero exactly what a contract asks you to zero, and measure where the zero cannot reach.

TODO 2.2 — the independent witness

J = torch.func.vmap(torch.func.jacrev(single))(z)  # B,d,d
sign, logabsdet = torch.linalg.slogdet(J)           # (B,)

single sends one point (d,) to one point (d,). Accept under atol = 1e-10, rtol = 1e-8 on the largest log-det; assert sign > 0.

Measured, 64 points: gap \(5.3 \times 10^{-16}\) (tol \(3.0 \times 10^{-10}\)) untrained; \(5.1 \times 10^{-15}\) (tol \(2.1 \times 10^{-8}\)) after training, where \(\max|\log|\det J|| = 2.11\).

Live-code the two lines. Say what each does: jacrev differentiates the single-example map; vmap batches it; slogdet reads the number the sum must match. The pitfall is the next slide, and it is worth the minute.

The pitfall: one determinant of the batch

torch.autograd.functional.jacobian(f, z) on the whole batch returns the Jacobian of the batched map \(\mathbb{R}^{Bd} \to \mathbb{R}^{Bd}\), shape (B, d, B, d).

Measured, \(B = 64\):

Quantity	Value
largest off-diagonal block entry	\(0.0\)
slogdet of the batched map	\(0.24336826533634\)
sum of the 64 per-example log-dets	\(0.24336826533634\)

One number for the whole batch. It cannot say which example is wrong; its plain det is a product of \(B\) factors and overflows for a trained flow; it costs \(B\) times the memory.

Block-diagonal, so slogdet equals the sum to 4e-16 — that is exactly why it is useless as a per-example check. The per-example stack is the object; vmap is how you get it.

TODO 2.3 — both round trips

Per example, float64:

• inverse(forward(z)) \(\approx\) z
• forward(inverse(x)) \(\approx\) x
• logdet_fwd + logdet_inv \(\approx 0\) at matching points

Measured, untrained: \(8.9 \times 10^{-16}\), \(8.9 \times 10^{-16}\), \(2.8 \times 10^{-17}\). Trained: \(9.1 \times 10^{-15}\), \(2.7 \times 10^{-15}\), \(1.0 \times 10^{-15}\).

The third line is U1.T1’s Exercise 3, executed: the inverse pass reads the same sum with the opposite sign.

Fast. Three asserts, one print. Tick object 1 on the board: a bijection, verified.

04Stack, and train

TODO 3.1 — compose, and the sign

Flow.forward: layers in order, log-dets add. Flow.inverse: layers reversed, each through its own inverse, log-dets add. Masks alternate.

\[ \log p_\theta(x) \;=\; \log p_0(z) \;+\; \log\big|\det J_{\psi^{-1}}(x)\big| , \qquad z = \psi^{-1}(x) \]

In code: z, logdet_inv = inverse(x); return log_p0(z) + logdet_inv. That is U1.T1’s \(\log p_0(z_0) - \sum_k \log|\det J_{\psi_k}(z_{k-1})|\) with the minus sign absorbed into logdet_inv. Written next to the implementation, so the sign is never implicit.

Thirty-eight to fifty-six minutes for block D. BOARD 2, right half: the mask alternation as two rows of A/B, and the sign line. Checkpoint 3.1 asserts log_prob == log_p0(z) − logdet_fwd(z): checked, not believed.

A permutation is free — in code

U1.T1’s Exercise 1 asked you to check in your head that a coordinate swap contributes \(0\) and undoes itself. Checkpoint 3.1 checks it in code:

mixed = Flow([layer_1, Permute([1, 0]), layer_2], 2)
y, ld = mixed(z)          # the swap adds exactly 0.0
z_back, _ = mixed.inverse(y)
assert (z_back - z).abs().max() < 1e-12

Alternating masks and inserted permutations are two spellings of the same idea: after two layers every coordinate has moved.

Note the subtlety in the notebook comment: the swap changes which half layer 2 reads, so the two-layer log-det changes; what must hold is that the swap ITSELF contributed zero and undoes itself. Also on this checkpoint: the witness passes on the whole six-layer stack.

TODO 3.2 — the closure is all you write

\[ \mathcal{L}(\theta) = -\,\mathbb{E}_{x_1 \sim q}\big[\log p_\theta(x_1)\big] \approx -\frac{1}{B}\sum_{i=1}^{B} \log p_\theta\big(x_1^{(i)}\big) \]

class LogProbModel(nn.Module):   # Trainer calls model(x)
    def forward(self, x):
        return self.flow.log_prob(x)

class NLL(nn.Module):            # loss_fn(model(x), y)
    def forward(self, logp, y):
        return -logp.mean()

Everything else is the given Trainer.

The subscript on the expectation is mandatory and it says: data. Selection on the validation split, never on a test set (there is none).

Settings and budget

• Optimizer, schedule — the course defaults: AdamW at \(3 \times 10^{-4}\), weight decay \(0.01\); warmup over 5% of the steps, then cosine, per step.
• Model — \(K = 6\) affine coupling layers, conditioners of width 64, \(s_{\max} = 2\): 26,508 parameters.
• Budget — 500 epochs \(\times\) 16 steps = 8000 steps at batch 256. Measured: 69 s on a laptop core.
• Data — train two_moons(4096, 0), validation two_moons(1024, 1). Selection on validation; there is no test set.

Deviation from the card is permitted in writing with a reason; none was needed. Halve the epochs on a slow machine — the checkpoint is relative.

The relative checkpoint

Pass if: the validation NLL at the end beats the first epoch’s by more than one nat, and the trained stack still passes the witness and the round trips.

Run	Time	Val NLL, epoch 1	Epoch 500
two-moons	69 s	4.674	2.099
gapped checkerboard	69 s	6.927	3.387

No absolute threshold: an absolute number is a promise about your hardware that nobody validated. TODO 3.3: same code, same budget, the checkerboard — one line changes.

Improvement 2.58 and 3.54 nats against a bar of 1. Run the checkerboard training while you talk through the next slide; keep the model, block E draws it.

Checkpoint — the curves

Two-moons settles near 2.10. The checkerboard settles near 3.39, half a nat above the dashed line: log 18 = 2.890 is the NLL of q itself, a uniform on eight squares of side 1.5. That half nat is measured and NOT explained today. For two-moons no closed-form entropy exists, so no line.

05Look at what you built

Two-moons: three views

Fifty-six to sixty-eight minutes for block E. Left: data gray, model samples orange. Middle: log_prob on a 200×200 grid, mass on the extent 0.9999 — exact density evaluation, the thing this unit exists for. Right: the source grid, [−3, 3]², pushed through ψ. TODO 4.1 draws the heatmap, 4.2 the grid; both are one call to a toy2d helper.

The warped grid is T1’s picture

• U1.T1 pushed one small cube through a diffeomorphism and read the density change off the volume change.
• This panel pushes the whole source grid through the learned diffeomorphism; the same reading holds cell by cell: large cell, density divided; crushed cell, density multiplied.
• The grid is one connected sheet, folded and stretched, never torn. A diffeomorphism cannot tear.

Keep that last sentence for the next figure.

Say “cannot tear” once, plainly, and stop. Do not continue the thought; U1.T2 does.

The source grid, folded into the moons

The right panel of the previous figure is the END of six layers. This is the six summands, on the trained weights, with the 400 gray source samples riding the sheet and one orange sample tracked with its running log|det J|.

Three things to say as it plays:

1. Each layer moves ONE coordinate; the dashed guide is the line the tracked sample is allowed to move along. The roles swap every layer — U1.T1’s alternation, on a model they trained.
2. The lines run off the panel and come back: after layer 1 more than half of the source grid is outside the window. That is what six coupling layers do to a 3σ box; the panel is a viewport, not a lie.
3. The counter is the log-det of the map on screen, not a tween. At the end the telescoped sum and the per-example autodiff Jacobian both read −1.105, and they agree to 8e-16. Under them, the equality of block 3 with the numbers filled in: −1.737 = −2.842 − (−1.105).

“Stretched and folded, never torn” is the caption of the final frame. Say it once more and stop; U1.T2 owns what comes next.

The gapped checkerboard: three views

Same three panels, support squares outlined. Let them look for ten seconds before asking the question on the next slide.

The question, and only the question

Important

What density does the model place in the hard-zero gaps, and what happens to the grid near those regions?

Measured on the reference model, \(400 \times 400\) grid:

• \(29.8\%\) of the model’s mass lies outside the eight squares, where \(q\) is exactly zero;
• along the corner-to-corner segment between two diagonal neighbours the density never drops below \(0.040\) (uniform level inside a square: \(1/18 \approx 0.056\)).

TODO 4.3 asks for the observation in words and in two numbers (mass outside; minimum on the segment). Let them compute both before showing the next slide.

The bridges appear, and stay

The 29.8% of the previous slide, as a curve in time. The reference run at sixteen checkpoints, epoch 1 to 500; the density on the figure’s own 400×400 grid, one colour scale for every frame; the validation NLL against log 18; the mass outside the eight squares, traced in the inset.

Three moments:

1. Epochs 1 to 30: one blurred blob spreads over the board, and the mass outside barely moves — 71% at epoch 1, 71% at epoch 30.
2. Epochs 30 to 50: the squares appear, and the mass outside drops 71% → 54%. The checkpoint at epoch 40 exists so the cross-fade never shows a state the model did not have.
3. From epoch 200 the curve bends flat: 33% → 30% over the last 300 epochs.

“And stay” is a measured claim: the precompute asserts the plateau two ways, one of them at epochs 200–300 while the learning rate is still above a third of its peak, so the cosine schedule alone cannot explain the flat tail.

The scene ends on the question, verbatim, and on “explained in U1.T2”. Say nothing about WHY. Next slide.

The observation, and where it goes

• The grid is stretched across every gap — one connected sheet, eight bulges, material in between.
• The validation NLL stops \(0.50\) nats above \(\log 18\), the NLL of \(q\) itself.

Explained in U1.T2. The solution records the observation and says exactly that. Do not read the size of the bridge as a theorem: what is exactly impossible for this model, and what is an effect of finite capacity or of optimization, are separated in that session.

“Explained in U1.T2” and nothing more — not even a hint. Tick nothing on the board; write the two numbers under a question mark.

06The library cross-check

Board 1 — two costs, two dialects

course \(\psi \colon z \to x\) · library \(f = \psi^{-1} \colon x \to z\)

Operation	Direction of \(\psi^{\theta}\)	Needs, per layer
Sampling	forward, \(x_0 \to x_1\)	\(\psi_k\)
Density / training	inverse, \(x_1 \to x_0\)	\(\psi_k^{-1}\) and \(\log\lvert\det J_{\psi_k}\rvert\)

Zuko exposes the density transform as its primary object. Never flip the course convention to match a library; write \(f = \psi^{-1}\) and translate.

Sixty-eight to eighty-two minutes for block F. BOARD 1 is U1.T1’s two-costs table with the direction line written above it. Both rows cost one conditioner pass in a coupling layer, in both dialects.

Same Trainer, one token changes

Zuko flows are nn.Modules whose call returns a distribution: flow() is \(p_\theta\).

zflow = zuko.flows.RealNVP(features=2, transforms=6,
                           hidden_features=(64, 64))

class ZukoLogProb(nn.Module):
    def forward(self, x):
        return self.flow().log_prob(x)  # flow(), then

Same data, same seeds, same optimizer, same schedule, same Trainer, same budget. That is the point of a cross-check.

TODO 5.1 is that one line. While it trains (74 s), read the direction table on the next slide.

The direction, asserted

	Map	Direction	Used by
course	\(\psi\)	\(z \to x\)	sample, your forward
flow.transform()	\(f = \psi^{-1}\)	\(x \to z\)	log_prob
flow.transform().inv	\(\psi\)	\(z \to x\)	sample

f.log_abs_det_jacobian(x, z) is \(\log|\det J_{\psi^{-1}}(x)|\) — your logdet_inv. Checked, not trusted: log_p0(f(x)) + ladj against flow().log_prob(x) on 1024 points, gap \(9.5 \times 10^{-7}\) (float32); f.inv(f(x)) returns \(x\) to \(3.3 \times 10^{-6}\).

Two asserts in the checkpoint cell. The sentence to land: the library’s forward is our inverse, and the arithmetic is identical — only the names are swapped.

The cross-check, passed

Scratch 2.099, Zuko 2.101 on the same validation points: 0.002 nats apart. Same parameter count, 26,508, because it is the same architecture; the differences are SiLU vs ReLU and the shape of the scale bound. Tick object 2 on the board: exact density evaluation, agreed by two implementations.

TODO 5.2 — source reading: two classes, one layer

Zuko splits the layer in two:

• the lazy class, GeneralCouplingTransform — owns the mask and the conditioner (the hypernetwork), builds a transform on demand;
• the transform, CouplingTransform — does the arithmetic, and delegates the per-coordinate map to the affine map, MonotonicAffineTransform.

Your one class maps onto the three of them. The next slide says what to find in each.

Open the files on the projector: import zuko; zuko.__file__ gives the directory. The notebook’s worksheet has the file:line anchors for zuko 1.5.0.

TODO 5.2 — four things to find

Your code	In Zuko	Find
mask	lazy: __init__; transform: split, merge	which value is the untouched half; how NICE alternates it
_heads	lazy: meta, hyper	what meta returns
inverse	transform and affine map: _inverse	which of _call / _inverse is \(\psi\)
the two log-dets	log_abs_det_jacobian; call_and_ladj	where the sum is taken; how the log-scale is bounded

Line anchors live in the notebook, generated from the pinned zuko 1.5.0 at build time. TODO 5.3 is reserved (only if U1.L2 is cut).

Answers in outline, for the close of the block: the mask is a buffer, ones mark the untouched half; meta IS _heads, returning (s, t) packaged as a transform; the library’s _call is our inverse and its _inverse is our forward, same arithmetic, names swapped; the sum over the moved coordinates happens one level up in DependentTransform; the bound is a soft clip at |s| < 6.9.

07JAX, and the close

The same map in JAX

def psi(z):                        # ONE example
    x = z
    for first_is_a, ws, s_cap in params:
        ia, ib = (0, 1) if first_is_a else (1, 0)
        raw = mlp(ws, x[ia:ia + 1])
        s = s_cap * jnp.tanh(raw[0] / s_cap)
        x = x.at[ib].set(x[ib] * jnp.exp(s) + raw[1])
    return x

J = jax.vmap(jax.jacfwd(psi))(z)   # (B, 2, 2)
_, logabsdet = jnp.linalg.slogdet(J)

Trained torch weights, copied over. Agreement with torch on 64 points: log-dets to \(1.1 \times 10^{-13}\), forward maps to \(3.9 \times 10^{-15}\).

Eighty-two to ninety minutes for block G, and this slide is skippable. jacfwd is forward mode, natural for a square map; vmap batches it, as torch.func.vmap did. A second language is a second independent witness. Tick object 3.

Three objects you own

Object	Checked by
a bijection: AffineCoupling + Flow	both round trips, per example, before and after training
exact density evaluation: Flow.log_prob	the sign check; the training curves; a library within \(0.002\) nats
an independent Jacobian witness	torch.func and JAX, per example, float64

And one thing you do not own yet: an explanation of the checkerboard’s gaps. \(29.8\%\) of the mass outside the support; \(0.50\) nats above \(\log 18\). U1.T2 opens there.

Read the board: three ticks and one open observation. Do not explain the fourth row.

What comes next

• Next session (U1.T2): it opens with today’s checkerboard and asks what a flow built from diffeomorphisms can and cannot represent exactly; then the rest of the zoo, and the depth question U1.T1 left open.
• Before then: finish TODO 4.3 in words, and use the notebook’s last cell to check your PS1.4(b) — the reference agrees with Flow.log_prob to \(1.8 \times 10^{-15}\).

Sources for today: (Dinh, Sohl-Dickstein, and Bengio 2017, secs. 3.2–3.3, §3.5), (Papamakarios et al. 2021, sec. 3.1.2), (Rozet et al. 2022).

Two sentences and out. PS1.4 was assigned in U1.T1; today’s module verifies part (b) numerically, and the disagreement, if any, names the intermediate point read by the wrong conditioner.

References

Dinh, Laurent, Jascha Sohl-Dickstein, and Samy Bengio. 2017. “Density Estimation Using Real NVP.” In International Conference on Learning Representations. https://arxiv.org/abs/1605.08803.

Papamakarios, George, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. 2021. “Normalizing Flows for Probabilistic Modeling and Inference.” Journal of Machine Learning Research 22 (57): 1–64. https://arxiv.org/abs/1912.02762.

Rozet, François et al. 2022. “Zuko: Normalizing Flows in PyTorch.” https://doi.org/10.5281/zenodo.7625672.