<jupyter_start><jupyter_text>Only OT for a Test<jupyter_code>n_epochs = 2001
eps = 1e-7
for i_epoch in range(n_epochs):
    point_dists = [th.distributions.MultivariateNormal(m, covariance_matrix=th.diag(th.exp(s) * th.exp(s)))
                   for m, s in zip(means_per_point, log_stds_per_point)]

    out_samples = class_dist.get_samples(0,100,)
    
    mixture_samples = th.cat([(th.randn(2,len(m)) * th.exp(l).unsqueeze(0)) + m.unsqueeze(0)
                          for m,l in zip(means_per_point, log_stds_per_point)])
    in_samples_mixture = mixture_model(mixture_samples)
    out_mixture_samples = feature_model_a(in_samples_mixture)
    
    all_out_samples = th.cat((out_samples, out_mixture_samples))

    prior_log_probs = class_dist.get_total_log_prob(0,all_out_samples)

    inverted = invert(feature_model_a, out_samples)
    inverted_to_mixture = invert(mixture_model, inverted)
    all_pre_in_samples = th.cat((inverted_to_mixture, mixture_samples))


    # jetzt likelihood inverted point ausrechnen
    log_probs_per_dist = th.stack([dist.log_prob(all_pre_in_samples) for dist in point_dists], 0)

    mixture_probs = th.mean(th.exp(log_probs_per_dist) + eps, dim=0)
    mixture_log_probs = th.log(mixture_probs)
    prior_probs = th.exp(prior_log_probs)
    sym_kl_div = -0.5 * th.sum(prior_probs * (mixture_log_probs - prior_log_probs)) - (
        0.5 * th.sum(mixture_probs * (prior_log_probs - mixture_log_probs)))
    in_diffs = in_samples_mixture.view(len(train_inputs), -1, train_inputs.shape[1]) - train_inputs.unsqueeze(1)
    OT_in = th.mean(th.norm(in_diffs,p=2, dim=-1))
    out_real = feature_model_a(train_inputs)
    out_diffs = out_mixture_samples.view(len(train_inputs), -1, out_real.shape[1]) - out_real.unsqueeze(1)
    OT_out = th.mean(th.norm(out_diffs,p=2, dim=-1))
    
    loss =  OT_in * 100 + OT_out * 100 #+ sym_kl_div
    optim_mixture_dist.zero_grad()
    optim_model_a.zero_grad()
    optim_dist.zero_grad()
    optim_mixture_model.zero_grad()
    loss.backward()
    optim_mixture_dist.step()
    optim_model_a.step()
    optim_dist.step()
    optim_mixture_model.step()
    if i_epoch % (n_epochs // 20) == 0:
        print("Epoch {:d} of {:d}".format(i_epoch, n_epochs))
        print("JS: {:.1E}".format(sym_kl_div.item()))
        print("OT: {:.1E}".format((OT_in + OT_out).item()))
        print("Mean std {:.1E}".format(th.mean(th.stack([th.exp(l) for l in log_stds_per_point]))))
        print("mean and std prior",class_dist.get_mean_std(0))
        

        with th. no_grad():
            out_samples = class_dist.get_samples(0,100,)
            inverted = invert(feature_model_a, out_samples)
            mixture_samples = th.cat([dist.sample((2,)) for dist in point_dists])
            in_mixture_samples = mixture_model(mixture_samples)
            out_mixture_samples = feature_model_a(in_mixture_samples)
            out_real = feature_model_a(train_inputs)
            
            # regenerate  transformed circles
            radians = np.linspace(0,2*np.pi,24)
            circle_points = np.stack([np.cos(radians), np.sin(radians)], axis=-1)
            circle_th = np_to_var(circle_points, device=train_inputs.device, dtype=np.float32)
            ms = th.stack(means_per_point)
            stds = th.exp(th.stack(log_stds_per_point))
            circles_per_point = ms.unsqueeze(1) + (circle_th.unsqueeze(0) * stds.unsqueeze(1))
            in_circles = mixture_model(circles_per_point.view(-1, circles_per_point.shape[-1]))
            out_circles = feature_model_a(in_circles)
            in_circles = in_circles.view(circles_per_point.shape)
            out_circles= out_circles.view(circles_per_point.shape)


        fig, axes = plt.subplots(1,2, figsize=(8,4))
        axes[0].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1],
                       alpha=0.75)
        axes[0].scatter(var_to_np(in_mixture_samples)[:,0], var_to_np(in_mixture_samples)[:,1],
                       alpha=0.75)
        axes[0].scatter(var_to_np(train_inputs)[:,0], var_to_np(train_inputs)[:,1])
        for c in var_to_np(in_circles):
            axes[0].plot(c[:,0], c[:,1],color='black')
        axes[0].set_title("Input space")
        axes[0].axis('equal')
        axes[1].scatter(var_to_np(out_samples)[:,0], var_to_np(out_samples)[:,1],
                       alpha=0.75,
                       label='Latent Prior')
        axes[1].scatter(var_to_np(out_mixture_samples)[:,0], var_to_np(out_mixture_samples)[:,1],
                       alpha=0.75,
                       label='Mixture')
        axes[1].scatter(var_to_np(out_real)[:,0], var_to_np(out_real)[:,1], label='Real data')
        for c in var_to_np(out_circles):
            axes[1].plot(c[:,0], c[:,1],color='black')
        axes[1].set_title("Output space")
        axes[1].axis('equal')
        axes[1].legend(bbox_to_anchor=(1,1,0,0))
        display_close(fig)<jupyter_output><empty_output><jupyter_text>Only KL-DIV<jupyter_code>n_epochs = 2001
eps = 1e-7
for i_epoch in range(n_epochs):
    point_dists = [th.distributions.MultivariateNormal(m, covariance_matrix=th.diag(th.exp(s) * th.exp(s)))
                   for m, s in zip(means_per_point, log_stds_per_point)]

    out_samples = class_dist.get_samples(0,100,)
    
    mixture_samples = th.cat([(th.randn(2,len(m)) * th.exp(l).unsqueeze(0)) + m.unsqueeze(0)
                          for m,l in zip(means_per_point, log_stds_per_point)])
    in_samples_mixture = mixture_model(mixture_samples)
    out_mixture_samples = feature_model_a(in_samples_mixture)
    
    all_out_samples = th.cat((out_samples, out_mixture_samples))

    prior_log_probs = class_dist.get_total_log_prob(0,all_out_samples)

    inverted = invert(feature_model_a, out_samples)
    inverted_to_mixture = invert(mixture_model, inverted)
    all_pre_in_samples = th.cat((inverted_to_mixture, mixture_samples))


    # jetzt likelihood inverted point ausrechnen
    log_probs_per_dist = th.stack([dist.log_prob(all_pre_in_samples) for dist in point_dists], 0)

    mixture_probs = th.mean(th.exp(log_probs_per_dist) + eps, dim=0)
    mixture_log_probs = th.log(mixture_probs)
    prior_probs = th.exp(prior_log_probs)
    sym_kl_div = -0.5 * th.sum(prior_probs * (mixture_log_probs - prior_log_probs)) - (
        0.5 * th.sum(mixture_probs * (prior_log_probs - mixture_log_probs)))
    in_diffs = in_samples_mixture.view(len(train_inputs), -1, train_inputs.shape[1]) - train_inputs.unsqueeze(1)
    OT_in = th.mean(th.norm(in_diffs,p=2, dim=-1))
    out_real = feature_model_a(train_inputs)
    out_diffs = out_mixture_samples.view(len(train_inputs), -1, out_real.shape[1]) - out_real.unsqueeze(1)
    OT_out = th.mean(th.norm(out_diffs,p=2, dim=-1))
    
    loss =   sym_kl_div + OT_in * 100 + OT_out * 100 
    optim_mixture_dist.zero_grad()
    optim_model_a.zero_grad()
    optim_dist.zero_grad()
    optim_mixture_model.zero_grad()
    loss.backward()
    optim_mixture_dist.step()
    optim_model_a.step()
    optim_dist.step()
    optim_mixture_model.step()
    if i_epoch % (n_epochs // 20) == 0:
        print("Epoch {:d} of {:d}".format(i_epoch, n_epochs))
        print("JS: {:.1E}".format(sym_kl_div.item()))
        print("OT: {:.1E}".format((OT_in + OT_out).item()))
        print("Mean std {:.1E}".format(th.mean(th.stack([th.exp(l) for l in log_stds_per_point]))))
        print("mean and std prior",class_dist.get_mean_std(0))
        

        with th. no_grad():
            out_samples = class_dist.get_samples(0,100,)
            inverted = invert(feature_model_a, out_samples)
            mixture_samples = th.cat([dist.sample((2,)) for dist in point_dists])
            in_mixture_samples = mixture_model(mixture_samples)
            out_mixture_samples = feature_model_a(in_mixture_samples)
            out_real = feature_model_a(train_inputs)
            
            # regenerate  transformed circles
            radians = np.linspace(0,2*np.pi,24)
            circle_points = np.stack([np.cos(radians), np.sin(radians)], axis=-1)
            circle_th = np_to_var(circle_points, device=train_inputs.device, dtype=np.float32)
            ms = th.stack(means_per_point)
            stds = th.exp(th.stack(log_stds_per_point))
            circles_per_point = ms.unsqueeze(1) + (circle_th.unsqueeze(0) * stds.unsqueeze(1))
            in_circles = mixture_model(circles_per_point.view(-1, circles_per_point.shape[-1]))
            out_circles = feature_model_a(in_circles)
            in_circles = in_circles.view(circles_per_point.shape)
            out_circles= out_circles.view(circles_per_point.shape)


        fig, axes = plt.subplots(1,2, figsize=(8,4))
        axes[0].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1],
                       alpha=0.75)
        axes[0].scatter(var_to_np(in_mixture_samples)[:,0], var_to_np(in_mixture_samples)[:,1],
                       alpha=0.75)
        axes[0].scatter(var_to_np(train_inputs)[:,0], var_to_np(train_inputs)[:,1])
        for c in var_to_np(in_circles):
            axes[0].plot(c[:,0], c[:,1],color='black')
        axes[0].set_title("Input space")
        axes[0].axis('equal')
        axes[1].scatter(var_to_np(out_samples)[:,0], var_to_np(out_samples)[:,1],
                       alpha=0.75,
                       label='Latent Prior')
        axes[1].scatter(var_to_np(out_mixture_samples)[:,0], var_to_np(out_mixture_samples)[:,1],
                       alpha=0.75,
                       label='Mixture')
        axes[1].scatter(var_to_np(out_real)[:,0], var_to_np(out_real)[:,1], label='Real data')
        for c in var_to_np(out_circles):
            axes[1].plot(c[:,0], c[:,1],color='black')
        axes[1].set_title("Output space")
        axes[1].axis('equal')
        axes[1].legend(bbox_to_anchor=(1,1,0,0))
        display_close(fig)
        


n_epochs = 20001
eps = 1e-7
for i_epoch in range(n_epochs):
    point_dists = [th.distributions.MultivariateNormal(m, covariance_matrix=th.diag(th.exp(s) * th.exp(s)))
                   for m, s in zip(means_per_point, log_stds_per_point)]

    out_samples = class_dist.get_samples(0,100,)
    
    mixture_samples = th.cat([(th.randn(2,len(m)) * th.exp(l).unsqueeze(0)) + m.unsqueeze(0)
                          for m,l in zip(means_per_point, log_stds_per_point)])
    in_samples_mixture = mixture_model(mixture_samples)
    out_mixture_samples = feature_model_a(in_samples_mixture)
    
    all_out_samples = th.cat((out_samples, out_mixture_samples))

    prior_log_probs = class_dist.get_total_log_prob(0,all_out_samples)

    inverted = invert(feature_model_a, out_samples)
    inverted_to_mixture = invert(mixture_model, inverted)
    all_pre_in_samples = th.cat((inverted_to_mixture, mixture_samples))


    # jetzt likelihood inverted point ausrechnen
    log_probs_per_dist = th.stack([dist.log_prob(all_pre_in_samples) for dist in point_dists], 0)

    mixture_probs = th.mean(th.exp(log_probs_per_dist) + eps, dim=0)
    mixture_log_probs = th.log(mixture_probs)
    prior_probs = th.exp(prior_log_probs)
    sym_kl_div = -0.5 * th.sum(prior_probs * (mixture_log_probs - prior_log_probs)) - (
        0.5 * th.sum(mixture_probs * (prior_log_probs - mixture_log_probs)))
    in_diffs = in_samples_mixture.view(len(train_inputs), -1, train_inputs.shape[1]) - train_inputs.unsqueeze(1)
    OT_in = th.mean(th.norm(in_diffs,p=2, dim=-1))
    out_real = feature_model_a(train_inputs)
    out_diffs = out_mixture_samples.view(len(train_inputs), -1, out_real.shape[1]) - out_real.unsqueeze(1)
    OT_out = th.mean(th.norm(out_diffs,p=2, dim=-1))
    
    loss =   sym_kl_div + OT_in * 100 + OT_out * 100 
    optim_mixture_dist.zero_grad()
    optim_model_a.zero_grad()
    optim_dist.zero_grad()
    optim_mixture_model.zero_grad()
    loss.backward()
    optim_mixture_dist.step()
    optim_model_a.step()
    optim_dist.step()
    optim_mixture_model.step()
    if i_epoch % (n_epochs // 20) == 0:
        print("Epoch {:d} of {:d}".format(i_epoch, n_epochs))
        print("JS: {:.1E}".format(sym_kl_div.item()))
        print("OT: {:.1E}".format((OT_in + OT_out).item()))
        print("Mean std {:.1E}".format(th.mean(th.stack([th.exp(l) for l in log_stds_per_point]))))
        print("mean and std prior",class_dist.get_mean_std(0))
        

        with th. no_grad():
            out_samples = class_dist.get_samples(0,100,)
            inverted = invert(feature_model_a, out_samples)
            mixture_samples = th.cat([dist.sample((2,)) for dist in point_dists])
            in_mixture_samples = mixture_model(mixture_samples)
            out_mixture_samples = feature_model_a(in_mixture_samples)
            out_real = feature_model_a(train_inputs)
            
            # regenerate  transformed circles
            radians = np.linspace(0,2*np.pi,24)
            circle_points = np.stack([np.cos(radians), np.sin(radians)], axis=-1)
            circle_th = np_to_var(circle_points, device=train_inputs.device, dtype=np.float32)
            ms = th.stack(means_per_point)
            stds = th.exp(th.stack(log_stds_per_point))
            circles_per_point = ms.unsqueeze(1) + (circle_th.unsqueeze(0) * stds.unsqueeze(1))
            in_circles = mixture_model(circles_per_point.view(-1, circles_per_point.shape[-1]))
            out_circles = feature_model_a(in_circles)
            in_circles = in_circles.view(circles_per_point.shape)
            out_circles= out_circles.view(circles_per_point.shape)


        fig, axes = plt.subplots(1,2, figsize=(8,4))
        axes[0].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1],
                       alpha=0.75)
        axes[0].scatter(var_to_np(in_mixture_samples)[:,0], var_to_np(in_mixture_samples)[:,1],
                       alpha=0.75)
        axes[0].scatter(var_to_np(train_inputs)[:,0], var_to_np(train_inputs)[:,1])
        for c in var_to_np(in_circles):
            axes[0].plot(c[:,0], c[:,1],color='black')
        axes[0].set_title("Input space")
        axes[0].axis('equal')
        axes[1].scatter(var_to_np(out_samples)[:,0], var_to_np(out_samples)[:,1],
                       alpha=0.75,
                       label='Latent Prior')
        axes[1].scatter(var_to_np(out_mixture_samples)[:,0], var_to_np(out_mixture_samples)[:,1],
                       alpha=0.75,
                       label='Mixture')
        axes[1].scatter(var_to_np(out_real)[:,0], var_to_np(out_real)[:,1], label='Real data')
        for c in var_to_np(out_circles):
            axes[1].plot(c[:,0], c[:,1],color='black')
        axes[1].set_title("Output space")
        axes[1].axis('equal')
        axes[1].legend(bbox_to_anchor=(1,1,0,0))
        display_close(fig)<jupyter_output><empty_output><jupyter_text>without learned prior model<jupyter_code>from reversible2.distribution import TwoClassDist

from reversible2.blocks import dense_add_block, conv_add_block_3x3
from reversible2.rfft import RFFT, Interleave
from reversible2.util import set_random_seeds
from torch.nn import ConstantPad2d
import torch as th
from reversible2.splitter import SubsampleSplitter


set_random_seeds(2019011641, cuda)
feature_model_a = nn.Sequential(
    dense_add_block(2,200),
    dense_add_block(2,200),
    dense_add_block(2,200),
    dense_add_block(2,200),
)
if cuda:
    feature_model_a.cuda()
from reversible2.ot_exact import ot_euclidean_loss_for_samples
class_dist = TwoClassDist(2,0, [0,1])
if cuda:
    class_dist.cuda()

means_per_point = [th.zeros(2,requires_grad=True) for _ in train_inputs]
log_stds_per_point = [th.zeros(2,requires_grad=True) for _ in train_inputs]

for m,l in zip(means_per_point, log_stds_per_point):
    m.data[:] = th.randn(2) * 0.05
    l.data[:] = -2
    
optim_model_a = th.optim.Adam(feature_model_a.parameters())
optim_dist = th.optim.Adam(class_dist.parameters(), lr=1e-2)
optim_mixture_dist = th.optim.Adam(means_per_point + log_stds_per_point, 
                                  lr=1e-2)
n_epochs = 20001
eps = 1e-7
for i_epoch in range(n_epochs):
    point_dists = [th.distributions.MultivariateNormal(m, covariance_matrix=th.diag(th.exp(s) * th.exp(s)))
                   for m, s in zip(means_per_point, log_stds_per_point)]

    out_samples = class_dist.get_samples(0,100,)
    
    mixture_samples = th.cat([(th.randn(2,len(m)) * th.exp(l).unsqueeze(0)) + m.unsqueeze(0)
                          for m,l in zip(means_per_point, log_stds_per_point)])
    out_mixture_samples = feature_model_a(mixture_samples)
    
    all_out_samples = th.cat((out_samples, out_mixture_samples))

    prior_log_probs = class_dist.get_total_log_prob(0,all_out_samples)

    inverted = invert(feature_model_a, out_samples)
    all_in_samples = th.cat((inverted, mixture_samples))


    # jetzt likelihood inverted point ausrechnen
    log_probs_per_dist = th.stack([dist.log_prob(all_in_samples) for dist in point_dists], 0)

    mixture_probs = th.mean(th.exp(log_probs_per_dist) + eps, dim=0)
    mixture_log_probs = th.log(mixture_probs)
    prior_probs = th.exp(prior_log_probs)
    sym_kl_div = -0.5 * th.sum(prior_probs * (mixture_log_probs - prior_log_probs)) - (
        0.5 * th.sum(mixture_probs * (prior_log_probs - mixture_log_probs)))
    in_diffs = mixture_samples.view(len(train_inputs), -1, train_inputs.shape[1]) - train_inputs.unsqueeze(1)
    OT_in = th.mean(th.norm(in_diffs,p=2, dim=-1))
    out_real = feature_model_a(train_inputs)
    out_diffs = out_mixture_samples.view(len(train_inputs), -1, out_real.shape[1]) - out_real.unsqueeze(1)
    OT_out = th.mean(th.norm(out_diffs,p=2, dim=-1))
    
    loss =  OT_in * 100 + OT_out * 100 + sym_kl_div
    optim_mixture_dist.zero_grad()
    optim_model_a.zero_grad()
    optim_dist.zero_grad()
    loss.backward()
    optim_mixture_dist.step()
    optim_model_a.step()
    optim_dist.step()
    if i_epoch % (n_epochs // 20) == 0:
        print("Epoch {:d} of {:d}".format(i_epoch, n_epochs))
        print("JS: {:.1E}".format(sym_kl_div.item()))
        print("OT: {:.1E}".format(OT.item()))
        print("Mean std {:.1E}".format(th.mean(th.stack([th.exp(l) for l in log_stds_per_point]))))
        print("mean and std prior",class_dist.get_mean_std(0))
        with th. no_grad():
            out_samples = class_dist.get_samples(0,100,)
            inverted = invert(feature_model_a, out_samples)
            mixture_samples = th.cat([dist.sample((2,)) for dist in point_dists])
            out_mixture_samples = feature_model_a(mixture_samples)
            out_real = feature_model_a(train_inputs)
        fig, axes = plt.subplots(1,2, figsize=(8,4))
        axes[0].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1])
        axes[0].scatter(var_to_np(mixture_samples)[:,0], var_to_np(mixture_samples)[:,1])
        axes[0].scatter(var_to_np(train_inputs)[:,0], var_to_np(train_inputs)[:,1])
        axes[0].set_title("Input space")
        axes[0].axis('equal')
        axes[1].scatter(var_to_np(out_samples)[:,0], var_to_np(out_samples)[:,1],
                       label='Latent Prior')
        axes[1].scatter(var_to_np(out_mixture_samples)[:,0], var_to_np(out_mixture_samples)[:,1],
                       label='Mixture')
        axes[1].scatter(var_to_np(out_real)[:,0], var_to_np(out_real)[:,1], label='Real data')
        axes[1].set_title("Output space")
        axes[1].axis('equal')
        axes[1].legend(bbox_to_anchor=(1,1,0,0))
        display_close(fig)
        


n_epochs = 2001
eps = 1e-7
for i_epoch in range(n_epochs):
    point_dists = [th.distributions.MultivariateNormal(m, covariance_matrix=th.diag(th.exp(s) * th.exp(s)))
    optim_mixt
                   for m, s in zip(means_per_point, log_stds_per_point)]

    out_samples = class_dist.get_samples(0,100,)
    
    mixture_samples = th.cat([(th.randn(2,len(m)) * th.exp(l).unsqueeze(0)) + m.unsqueeze(0)
                          for m,l in zip(means_per_point, log_stds_per_point)])
    out_mixture_samples = feature_model_a(mixture_samples)
    
    all_out_samples = th.cat((out_samples, out_mixture_samples))

    prior_log_probs = class_dist.get_total_log_prob(0,all_out_samples)

    inverted = invert(feature_model_a, out_samples)
    all_in_samples = th.cat((inverted, mixture_samples))


    # jetzt likelihood inverted point ausrechnen
    log_probs_per_dist = th.stack([dist.log_prob(all_in_samples) for dist in point_dists], 0)

    mixture_probs = th.mean(th.exp(log_probs_per_dist) + eps, dim=0)
    mixture_log_probs = th.log(mixture_probs)
    prior_probs = th.exp(prior_log_probs)
    sym_kl_div = -0.5 * th.sum(prior_probs * (mixture_log_probs - prior_log_probs)) - (
        0.5 * th.sum(mixture_probs * (prior_log_probs - mixture_log_probs)))
    diffs = mixture_samples.view(len(train_inputs), -1, mixture_samples.shape[1]) - train_inputs.unsqueeze(1)
    OT = th.mean(th.norm(diffs,p=2, dim=-1))
    loss =  OT * 100 + sym_kl_div
    optim_mixture_dist.zero_grad()
    optim_model_a.zero_grad()
    optim_dist.zero_grad()
    loss.backward()
    optim_mixture_dist.step()
    optim_model_a.step()
    optim_dist.step()
    if i_epoch % (n_epochs // 20) == 0:
        print("Epoch {:d} of {:d}".format(i_epoch, n_epochs))
        print("JS: {:.1E}".format(sym_kl_div.item()))
        print("OT: {:.1E}".format(OT.item()))
        print("Mean std {:.1E}".format(th.mean(th.stack([th.exp(l) for l in log_stds_per_point]))))
        print("mean and std prior",class_dist.get_mean_std(0))
        with th. no_grad():
            out_samples = class_dist.get_samples(0,100,)
            inverted = invert(feature_model_a, out_samples)
            mixture_samples = th.cat([dist.sample((2,)) for dist in point_dists])
            out_mixture_samples = feature_model_a(mixture_samples)
            out_real = feature_model_a(train_inputs)
        fig, axes = plt.subplots(1,2, figsize=(8,4))
        axes[0].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1])
        axes[0].scatter(var_to_np(mixture_samples)[:,0], var_to_np(mixture_samples)[:,1])
        axes[0].scatter(var_to_np(train_inputs)[:,0], var_to_np(train_inputs)[:,1])
        axes[0].set_title("Input space")
        axes[0].axis('equal')
        axes[1].scatter(var_to_np(out_samples)[:,0], var_to_np(out_samples)[:,1],
                       label='Latent Prior')
        axes[1].scatter(var_to_np(out_mixture_samples)[:,0], var_to_np(out_mixture_samples)[:,1],
                       label='Mixture')
        axes[1].scatter(var_to_np(out_real)[:,0], var_to_np(out_real)[:,1], label='Real data')
        axes[1].set_title("Output space")
        axes[1].axis('equal')
        axes[1].legend(bbox_to_anchor=(1,1,0,0))
        display_close(fig)
        


n_epochs = 500
eps = 1e-7
for i_epoch in range(n_epochs):
    point_dists = [th.distributions.MultivariateNormal(m, covariance_matrix=th.diag(th.exp(s) * th.exp(s)))
                   for m, s in zip(means_per_point, log_stds_per_point)]

    out_samples = class_dist.get_samples(0,100,)
    
    
    mixture_samples = th.cat([dist.sample((2,)) for dist in point_dists])
    out_mixture_samples = feature_model_a(mixture_samples)
    
    all_out_samples = th.cat((out_samples, out_mixture_samples))

    prior_log_probs = class_dist.get_total_log_prob(0,all_out_samples)

    inverted = invert(feature_model_a, out_samples)
    all_in_samples = th.cat((inverted, mixture_samples))


    # jetzt likelihood inverted point ausrechnen
    log_probs_per_dist = th.stack([dist.log_prob(all_in_samples) for dist in point_dists], 0)

    mixture_probs = th.mean(th.exp(log_probs_per_dist) + eps, dim=0)
    mixture_log_probs = th.log(mixture_probs)
    prior_probs = th.exp(prior_log_probs)
    sym_kl_div = -0.5 * th.sum(prior_probs * (mixture_log_probs - prior_log_probs)) - (
        0.5 * th.sum(mixture_probs * (prior_log_probs - mixture_log_probs)))
    optim_mixture_dist.zero_grad()
    optim_model_a.zero_grad()
    optim_dist.zero_grad()
    sym_kl_div.backward()
    optim_mixture_dist.step()
    optim_model_a.step()
    optim_dist.step()
    if i_epoch % (n_epochs // 20) == 0:
        print("{:.1E}".format(sym_kl_div.item()))
        print("Mean std {:.1E}".format(th.mean(th.stack([th.exp(l) for l in log_stds_per_point]))))
        with th. no_grad():
            out_samples = class_dist.get_samples(0,100,)
            inverted = invert(feature_model_a, out_samples)
            mixture_samples = th.cat([dist.sample((2,)) for dist in point_dists])
            out_mixture_samples = feature_model_a(mixture_samples)
        fig, axes = plt.subplots(1,2, figsize=(8,4))
        axes[0].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1])
        axes[0].scatter(var_to_np(mixture_samples)[:,0], var_to_np(mixture_samples)[:,1])
        axes[0].set_title("Input space")
        axes[1].scatter(var_to_np(out_samples)[:,0], var_to_np(out_samples)[:,1],
                       label='Latent Prior')
        axes[1].scatter(var_to_np(out_mixture_samples)[:,0], var_to_np(out_mixture_samples)[:,1],
                       label='Mixture')
        axes[1].set_title("Output space")
        axes[1].legend(bbox_to_anchor=(1,1,0,0))
        display_close(fig)<jupyter_output><empty_output><jupyter_text>before only look at samples from prior<jupyter_code>n_epochs = 500
eps = 1e-7
for i_epoch in range(n_epochs):
    point_dists = [th.distributions.MultivariateNormal(m, covariance_matrix=th.diag(th.exp(s) * th.exp(s)))
                   for m, s in zip(means_per_point, log_stds_per_point)]

    out_samples = class_dist.get_samples(0,100,)

    prior_log_probs = class_dist.get_total_log_prob(0,out_samples)

    inverted = invert(feature_model_a, out_samples)


    # jetzt likelihood inverted point ausrechnen
    log_probs_per_dist = th.stack([dist.log_prob(inverted) for dist in point_dists], 0)

    mixture_probs = th.mean(th.exp(log_probs_per_dist) + eps, dim=0)

    mixture_log_probs = th.log(mixture_probs)
    prior_probs = th.exp(prior_log_probs)
    sym_kl_div = -0.5 * th.sum(prior_probs * (mixture_log_probs - prior_log_probs)) - (
        0.5 * th.sum(mixture_probs * (prior_log_probs - mixture_log_probs)))
    optim_mixture_dist.zero_grad()
    optim_model_a.zero_grad()
    optim_dist.zero_grad()
    sym_kl_div.backward()
    optim_mixture_dist.step()
    optim_model_a.step()
    optim_dist.step()
    if i_epoch % (n_epochs // 20) == 0:
        print("{:.1E}".format(sym_kl_div.item()))
        print("Mean std {:.1E}".format(th.mean(th.stack([th.exp(l) for l in log_stds_per_point]))))
        with th. no_grad():
            out_samples = class_dist.get_samples(0,100,)
            mixture_samples = th.cat([dist.sample((2,)) for dist in point_dists])
            out_mixture_samples = feature_model_a(mixture_samples)
            inverted = invert(feature_model_a, out_samples)
        fig, axes = plt.subplots(1,2, figsize=(8,4))
        axes[0].scatter(var_to_np(inverted)[:,0], var_to_np(inverted)[:,1])
        axes[0].scatter(var_to_np(mixture_samples)[:,0], var_to_np(mixture_samples)[:,1])
        axes[0].set_title("Input space")
        axes[1].scatter(var_to_np(out_samples)[:,0], var_to_np(out_samples)[:,1],
                       label='Latent Prior')
        axes[1].scatter(var_to_np(out_mixture_samples)[:,0], var_to_np(out_mixture_samples)[:,1],
                       label='Mixture')
        axes[1].set_title("Output space")
        axes[1].legend(bbox_to_anchor=(1,1,0,0))
        display_close(fig)
        


class_dist.get_mean_std(0)
# associate a gaussian with every point
# bring it forward<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13801364890373166, 'fraction_numerical': 0.022832873529838826, 'mean_word_length': 3.389579349904398, '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': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '182971', 'license': 'MIT', 'n_tokens_mistral': 10773, 'n_tokens_neox': 10408, 'n_words': 1364, 'path': 'notebooks/toy-1d-2d-examples/MoonWithContinuous.ipynb', 'repo_name': 'robintibor/reversible2', 'size': '27548'}