<jupyter_start><jupyter_text>Let's generate some synthetic data for this example.<jupyter_code># Define trainX, valX, testX
import numpy as np
import torch

# Set seeds for reproducibility
np.random.seed(123) 
torch.random.manual_seed(123)

x = np.random.chisquare(10, size=1000)
y = x ** 2 + np.random.normal(size=1000) * 50
X = torch.Tensor(np.stack([x, y], 1))

trainX, valX, testX = X[:800], X[800:900], X[900:]

plt.scatter(*X.numpy().T, alpha=.1);<jupyter_output><empty_output><jupyter_text>Create the flow, train it, plot its loss evolution and compute the test loss.<jupyter_code>dim = trainX.size(1) # dimension of the flow

device = get_device() # cuda if available, cpu otherwise

# X transformed to base distribution U, like so:
# X -> BatchNorm -> MADE-DSF -> BatchNorm -> MADE-DSF -> U
flow = Sequential(
    BatchNorm(dim=dim), # we use BatchNorm for training stabilization
    MADE(DSF(dim=dim)), # combines a MADE conditioner with an Affine transformer
    
    BatchNorm(dim=dim),
    MADE(DSF(dim=dim))
).to(device) # don't forget to send it to device

# Train the flow with the train function, but you can use your own.
# train uses early stopping, that's why we need a validation set.
train_losses, val_losses = train(flow, trainX, valX)

# Plot training and validation losses
plot_losses(train_losses, val_losses)

# Compute the test set negative log-likelihood (the loss function used for training)
test_nll(flow, testX)<jupyter_output>690it [01:23,  8.28it/s, current_loss=8.42, best_epoch=590, best_loss=8.29]<jupyter_text>Plot some samples to see if they reflect the real distribution.<jupyter_code>with torch.no_grad():
    sample = flow.sample(trainX.size(0))
    sample = sample.cpu().numpy() # to numpy array
    
plt.scatter(*trainX.numpy().T, alpha=.1, label='real')
plt.scatter(*sample.T, alpha=.1, label='fake')
plt.legend();<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12652196929592377, 'fraction_numerical': 0.029645314981471677, 'mean_word_length': 5.038338658146965, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '389808', 'license': 'MIT', 'n_tokens_mistral': 705, 'n_tokens_neox': 650, 'n_words': 191, 'path': 'examples/README-example.ipynb', 'repo_name': 'aparafita/flow-torch', 'size': '1889'}