import torch
import zuko

# Create a flow with 3 sample features and 5 context features
flow = zuko.flows.NSF(3, 5, transforms=3, hidden_features=[128] * 3)

# Train to maximize log-likelihood
optimizer = torch.optim.Adam(flow.parameters(), lr=1e-3)

for x, c in trainset:
    loss = -flow(c).log_prob(x)  # -log p(x | c)
    loss = loss.mean()
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Sample from the learned distribution
x = flow(c_star).sample((64,))