<jupyter_start><jupyter_text>Optimal transport<jupyter_code>from reversible2.ot_exact import ot_euclidean_loss_for_samples

x_fake = np_to_var(np.linspace(-2,2, 400), dtype=np.float32, requires_grad=True)

zero_(valid_x.grad)
zero_(x_fake.grad)
for i in range(2):
    this_x = valid_x[i:i+1]
    ot_loss = ot_euclidean_loss_for_samples(this_x.unsqueeze(1), x_fake.unsqueeze(1))
    ot_loss.backward()
plt.figure(figsize=(8,3))
plt.plot(var_to_np(x_fake), var_to_np(x_fake.grad /2) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad / 2)))

zero_(valid_x.grad)
zero_(x_fake.grad)
ot_loss = ot_euclidean_loss_for_samples(valid_x.unsqueeze(1), x_fake.unsqueeze(1))
ot_loss.backward()
plt.plot(var_to_np(x_fake), var_to_np(x_fake.grad) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad)))
plt.legend(("Gradient Batched", "Integral Batched", "Gradient Full", "Integral Full"),
          bbox_to_anchor=(1,1,0,0))
from reversible2.ot_exact import ot_euclidean_loss_for_samples

x_fake = np_to_var(np.linspace(-2,2, 400), dtype=np.float32, requires_grad=True)
x_real = th.cat((x, valid_x)).detach().requires_grad_(True)
zero_(x_real.grad)
zero_(x_fake.grad)
for i in range(4):
    this_x = x_real[i:i+1]
    ot_loss = ot_euclidean_loss_for_samples(this_x.unsqueeze(1), x_fake.unsqueeze(1))
    ot_loss.backward()
plt.figure(figsize=(8,3))
plt.plot(var_to_np(x_fake), -var_to_np(x_fake.grad /4) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad / 4)))

zero_(x_real.grad)
zero_(x_fake.grad)
ot_loss = ot_euclidean_loss_for_samples(x_real.unsqueeze(1), x_fake.unsqueeze(1))
ot_loss.backward()
plt.plot(var_to_np(x_fake), -var_to_np(x_fake.grad) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad)))

plt.legend(("Gradient Batched", "Integral Batched", "Gradient Full", "Integral Full"),
          bbox_to_anchor=(1,1,0,0))
from reversible2.ot_exact import ot_euclidean_loss_for_samples

x_fake = np_to_var(np.linspace(-2,2, 400), dtype=np.float32, requires_grad=True)
x_real = th.cat((x, valid_x)).detach().requires_grad_(True)
zero_(x_real.grad)
zero_(x_fake.grad)

all_pairs = list(itertools.product(list(range(4)), list(range(4))))
for inds in all_pairs:
    this_x = x_real[np.array(inds).astype(np.int64)]
    ot_loss = ot_euclidean_loss_for_samples(this_x.unsqueeze(1), x_fake.unsqueeze(1))
    ot_loss.backward()
plt.figure(figsize=(8,3))
plt.plot(var_to_np(x_fake), -var_to_np(x_fake.grad /len(all_pairs)) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad / len(all_pairs))))

zero_(x_real.grad)
zero_(x_fake.grad)
ot_loss = ot_euclidean_loss_for_samples(x_real.unsqueeze(1), x_fake.unsqueeze(1))
ot_loss.backward()
plt.plot(var_to_np(x_fake), -var_to_np(x_fake.grad) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad)))

plt.legend(("Gradient Batched", "Integral Batched", "Gradient Full", "Integral Full"),
          bbox_to_anchor=(1,1,0,0))
from reversible2.ot_exact import ot_euclidean_loss_for_samples

x_fake = np_to_var(np.linspace(-2,2, 400), dtype=np.float32, requires_grad=True)
x_real = th.linspace(-1.5,1.5,6).detach().requires_grad_(True)

grads = []
zero_(x_real.grad)
zero_(x_fake.grad)
for i in range(len(x_real)):
    this_x = x_real[i:i+1]
    ot_loss = ot_euclidean_loss_for_samples(this_x.unsqueeze(1), x_fake.unsqueeze(1))
    ot_loss.backward()
plt.figure(figsize=(8,3))
#plt.plot(var_to_np(x_fake), -var_to_np(x_fake.grad /len(x_real)) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad / len(x_real))), label='Integral 1')
grads.append(-np.cumsum(var_to_np(x_fake.grad / len(x_real))))
zero_(x_real.grad)
zero_(x_fake.grad)

all_pairs = list(itertools.product(list(range(len(x_real))), list(range(len(x_real)))))
for inds in all_pairs:
    this_x = x_real[np.array(inds).astype(np.int64)]
    ot_loss = ot_euclidean_loss_for_samples(this_x.unsqueeze(1), x_fake.unsqueeze(1))
    ot_loss.backward()
#plt.plot(var_to_np(x_fake), -var_to_np(x_fake.grad /len(all_pairs)) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad / len(all_pairs))), label='Integral Batched')
grads.append(-np.cumsum(var_to_np(x_fake.grad / len(all_pairs))))

zero_(x_real.grad)
zero_(x_fake.grad)
ot_loss = ot_euclidean_loss_for_samples(x_real.unsqueeze(1), x_fake.unsqueeze(1))
ot_loss.backward()
#plt.plot(var_to_np(x_fake), -var_to_np(x_fake.grad) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad)), label='Integral Full')
grads.append(-np.cumsum(var_to_np(x_fake.grad)))

plt.plot(var_to_np(x_fake), np.mean(grads, axis=0), label='Integral mean')
plt.scatter(var_to_np(x_real), var_to_np(x_real) * 0)
plt.legend(bbox_to_anchor=(1,1,0,0))
from reversible2.ot_exact import ot_euclidean_loss_for_samples

x_fake = np_to_var(np.linspace(-2,2, 400), dtype=np.float32, requires_grad=True)
x_real = th.cat((x, valid_x)).detach().requires_grad_(True)


zero_(x_real.grad)
zero_(x_fake.grad)
for i in range(len(x_real)):
    this_x = x_real[i:i+1]
    XY = th.mean(th.abs(this_x.unsqueeze(1) - x_fake.unsqueeze(0)))
    XX = th.mean(th.abs(this_x.unsqueeze(1) - this_x.unsqueeze(0)))
    YY = th.mean(th.abs(x_fake.unsqueeze(1) - x_fake.unsqueeze(0)))
    ED = 2 * th.mean(XY) -th.mean(XX) - th.mean(YY)
    ED.backward()
plt.figure(figsize=(8,3))
#plt.plot(var_to_np(x_fake), -var_to_np(x_fake.grad /4) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad / len(x_real))),
        label='Integral 1')

zero_(x_real.grad)
zero_(x_fake.grad)
all_pairs = list(itertools.product(list(range(4)), list(range(4))))
for inds in all_pairs:
    this_x = x_real[np.array(inds).astype(np.int64)]
    XY = th.mean(th.abs(this_x.unsqueeze(1) - x_fake.unsqueeze(0)))
    XX = th.mean(th.abs(this_x.unsqueeze(1) - this_x.unsqueeze(0)))
    YY = th.mean(th.abs(x_fake.unsqueeze(1) - x_fake.unsqueeze(0)))
    ED = 2 * th.mean(XY) -th.mean(XX) - th.mean(YY)
    ED.backward()
#plt.plot(var_to_np(x_fake), -var_to_np(x_fake.grad /len(all_pairs)) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad / len(all_pairs))),
        label='Integral Batched')
zero_(x_real.grad)
zero_(x_fake.grad)
this_x = x_real
XY = th.mean(th.abs(this_x.unsqueeze(1) - x_fake.unsqueeze(0)))
XX = th.mean(th.abs(this_x.unsqueeze(1) - this_x.unsqueeze(0)))
YY = th.mean(th.abs(x_fake.unsqueeze(1) - x_fake.unsqueeze(0)))
ED = 2 * th.mean(XY) -th.mean(XX) - th.mean(YY)
ED.backward()
#plt.plot(var_to_np(x_fake), -var_to_np(x_fake.grad) * 100)
plt.plot(var_to_np(x_fake), -np.cumsum(var_to_np(x_fake.grad)),
        label='Integral Full')
plt.scatter(var_to_np(x_real), var_to_np(x_real) * 0)
plt.legend(bbox_to_anchor=(1,1,0,0))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.17363392057330546, 'fraction_numerical': 0.026276500447894893, 'mean_word_length': 6.451612903225806, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '597903', 'license': 'MIT', 'n_tokens_mistral': 3213, 'n_tokens_neox': 3043, 'n_words': 352, 'path': 'notebooks/toy-1d-2d-examples/regulate-norm/Energy_Distance_OT_Comparison.ipynb', 'repo_name': 'robintibor/reversible2', 'size': '6698'}