Normalizing Flows (RealNVP)

import math, torch, torch.nn as nn
from sklearn.datasets import make_moons

torch.manual_seed(0)
X = torch.tensor(make_moons(2000, noise=0.05, random_state=0)[0], dtype=torch.float32)

class Coupling(nn.Module):
    """RealNVP coupling layer: keep one coordinate, scale+shift the other."""
    def __init__(self, flip):
        super().__init__()
        self.flip = flip
        self.net = nn.Sequential(nn.Linear(1, 64), nn.ReLU(), nn.Linear(64, 2))
    def forward(self, x):                     # data -> latent
        a, b = (x[:, 1:], x[:, :1]) if self.flip else (x[:, :1], x[:, 1:])
        s, t = self.net(a).chunk(2, 1); s = torch.tanh(s)
        b = b * torch.exp(s) + t
        out = torch.cat([b, a], 1) if self.flip else torch.cat([a, b], 1)
        return out, s.sum(1)                  # log|det J| is just the sum of s
    def inverse(self, z):                     # latent -> data (exact inverse)
        a, b = (z[:, 1:], z[:, :1]) if self.flip else (z[:, :1], z[:, 1:])
        s, t = self.net(a).chunk(2, 1); s = torch.tanh(s)
        b = (b - t) * torch.exp(-s)
        return torch.cat([b, a], 1) if self.flip else torch.cat([a, b], 1)

layers = nn.ModuleList([Coupling(i % 2 == 1) for i in range(6)])
opt = torch.optim.Adam(layers.parameters(), lr=1e-3)

def log_prob(x):
    logdet = 0
    for layer in layers:
        x, ld = layer(x); logdet = logdet + ld
    return (-0.5 * (x ** 2).sum(1) - math.log(2 * math.pi)) + logdet   # change of variables

for step in range(1500):
    loss = -log_prob(X).mean()
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 500 == 0 or step == 1499:
        print(f"step {step:4d}  avg log-likelihood {-loss.item():.3f}")

with torch.no_grad():
    z = torch.randn(5, 2)
    x = z
    for layer in reversed(layers):
        x = layer.inverse(x)
    back = x
    for layer in layers:
        back, _ = layer(back)
print("generated points:", [[round(v, 2) for v in p] for p in x.tolist()])
print("invertible:", torch.allclose(back, z, atol=1e-4))