<jupyter_start><jupyter_text>Visualize training and testing data<jupyter_code>_, axes = plt.subplots(1,2,figsize=(12,4))
x,y = train_loader.dataset.array[:,0], train_loader.dataset.array[:,1]
_ = axes[0].scatter(x,y)
_ = axes[0].set_title('Training data')

x,y = test_loader.dataset.array[:,0], test_loader.dataset.array[:,1]
_ = axes[1].scatter(x,y)
_ = axes[1].set_title('Testing data')<jupyter_output><empty_output><jupyter_text>Train the flow model<jupyter_code>low = torch.FloatTensor([0, 0])
high = torch.FloatTensor([1, 1])
target_distribution = Uniform(low, high)
flow, train_losses, test_losses = train_and_eval(100, 5e-3, train_loader, test_loader, target_distribution)<jupyter_output><empty_output><jupyter_text>View the losses<jupyter_code>_ = plt.plot(train_losses, label='train_loss')
_ = plt.plot(test_losses, label='test_loss')
_ = plt.legend()
xs = np.linspace(-1.5,2.5,200)
ys = np.linspace(-1,1.5,200)
xs, ys = np.meshgrid(xs, ys)
x = np.stack([xs, ys], axis=2).reshape(-1,2)
with torch.no_grad():
    z, dz_by_dx = flow(torch.FloatTensor(x))
    px = (target_distribution.log_prob(z) + dz_by_dx.log())
    px = px.sum(1).exp()
    px = px.cpu().numpy().reshape(200,200)
    plt.pcolor(xs,ys,px)<jupyter_output><empty_output><jupyter_text>Empirical distirbution of `z`<jupyter_code>with torch.no_grad():
    z, _ = flow(torch.FloatTensor(train_loader.dataset.array))
    z = z.cpu().numpy()
    plt.scatter(z[:,0], z[:,1])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18218898708361658, 'fraction_numerical': 0.03195105370496261, 'mean_word_length': 6.666666666666667, 'pattern_counts': {'":': 0, '<': 17, '<?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': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '440849', 'license': 'MIT', 'n_tokens_mistral': 628, 'n_tokens_neox': 614, 'n_words': 80, 'path': '2d_autoregressive/2d.ipynb', 'repo_name': 'TinyVolt/normalizing-flows', 'size': '1471'}