stein.py

# -*- coding: utf-8 -*-
"""Stein.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1WiuK6JjFGMWqX3XqWkKZxOfHfMLvFAL1
"""

# Paste this to site-packages/torch/distributions/distribution.py (replace the constructor)

#     def __init__(self, batch_shape=torch.Size(), event_shape=torch.Size(), validate_args=None):
#         self._batch_shape = batch_shape
#         self._event_shape = event_shape
#         if validate_args is not None:
#             self._validate_args = validate_args
#         if self._validate_args:
#             try:
#                 arg_constraints = self.arg_constraints
#             except NotImplementedError:
#                 arg_constraints = {}
#                 warnings.warn(f'{self.__class__} does not define `arg_constraints`. ' +
#                               'Please set `arg_constraints = {}` or initialize the distribution ' +
#                               'with `validate_args=False` to turn off validation.')
#             for param, constraint in arg_constraints.items():
#                 if param in ["precision_matrix", 'scale_tril']:
#                     continue
#                 if constraints.is_dependent(constraint):
#                     continue  # skip constraints that cannot be checked
#                 if param not in self.__dict__ and isinstance(getattr(type(self), param), lazy_property):
#                     continue  # skip checking lazily-constructed args
#                 if not constraint.check(getattr(self, param)).all():
#                     raise ValueError("The parameter {} has invalid values".format(param))
#         super(Distribution, self).__init__()


import altair as alt
import numpy as np
import pandas as pd
import torch
import torch.autograd as autograd
import torch.optim as optim
from torch.autograd import Variable
from tqdm.auto import tqdm

alt.data_transformers.enable("default", max_rows=None)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


def get_density_chart(P, d=7.0, step=0.1):
    xv, yv = torch.meshgrid([torch.arange(-d, d, step), torch.arange(-d, d, step)])
    pos_xy = torch.cat((xv.unsqueeze(-1), yv.unsqueeze(-1)), dim=-1)
    p_xy = P.log_prob(pos_xy.to(device)).exp().unsqueeze(-1).cpu()
    #   print(p_xy)
    df = torch.cat([pos_xy, p_xy], dim=-1).numpy()
    df = pd.DataFrame(
        {
            "x": df[:, :, 0].ravel(),
            "y": df[:, :, 1].ravel(),
            "p": df[:, :, 2].ravel(),
        }
    )

    chart = (
        alt.Chart(df)
        .mark_point()
        .encode(
            x="x:Q",
            y="y:Q",
            color=alt.Color("p:Q", scale=alt.Scale(scheme="viridis")),
            tooltip=["x", "y", "p"],
        )
    )

    return chart


def get_density_charts(Ps, d=7.0, step=0.1):
    xv, yv = torch.meshgrid([torch.arange(-d, d, step), torch.arange(-d, d, step)])
    pos_xy = torch.cat((xv.unsqueeze(-1), yv.unsqueeze(-1)), dim=-1)

    #   print([P.log_prob(pos_xy).exp().unsqueeze(-1).cpu() for P in Ps])
    p_xy = torch.stack([P.log_prob(pos_xy.to(device)).exp().unsqueeze(-1).cpu() for P in Ps]).mean(0)

    df = torch.cat([pos_xy, p_xy], dim=-1).numpy()
    df = pd.DataFrame(
        {
            "x": df[:, :, 0].ravel(),
            "y": df[:, :, 1].ravel(),
            "p": df[:, :, 2].ravel(),
        }
    )

    chart = (
        alt.Chart(df)
        .mark_point()
        .encode(
            x="x:Q",
            y="y:Q",
            color=alt.Color("p:Q", scale=alt.Scale(scheme="viridis")),
            tooltip=["x", "y", "p"],
        )
    )

    return chart


def get_particles_chart(X):
    df = pd.DataFrame(
        {
            "x": X[:, 0],
            "y": X[:, 1],
        }
    )

    chart = alt.Chart(df).mark_circle(color="red").encode(x="x:Q", y="y:Q")

    return chart


class MinNormSolver:
    MAX_ITER = 250
    STOP_CRIT = 1e-5

    def _min_norm_element_from2(v1v1, v1v2, v2v2):
        if v1v2 >= v1v1:
            # Case: Fig 1, third column
            gamma = 0.999
            cost = v1v1
            return gamma, cost
        if v1v2 >= v2v2:
            # Case: Fig 1, first column
            gamma = 0.001
            cost = v2v2
            return gamma, cost
        # Case: Fig 1, second column
        gamma = -1.0 * ((v1v2 - v2v2) / (v1v1 + v2v2 - 2 * v1v2))
        cost = v2v2 + gamma * (v1v2 - v2v2)
        return gamma, cost

    def _min_norm_2d(vecs, dps, decomposition):

        if decomposition:
            dmin = 5e8
            for i in range(len(vecs)):
                for j in range(i + 1, len(vecs)):
                    if (i, j) not in dps:
                        dps[(i, j)] = 0.0
                        for k in range(len(vecs[i])):
                            dps[(i, j)] += torch.mul(vecs[i][k], vecs[j][k]).sum()
                        dps[(j, i)] = dps[(i, j)]
                    if (i, i) not in dps:
                        dps[(i, i)] = 0.0
                        for k in range(len(vecs[i])):
                            dps[(i, i)] += torch.mul(vecs[i][k], vecs[i][k]).sum()
                    if (j, j) not in dps:
                        dps[(j, j)] = 0.0
                        for k in range(len(vecs[i])):
                            dps[(j, j)] += torch.mul(vecs[j][k], vecs[j][k]).sum()
                    c, d = MinNormSolver._min_norm_element_from2(dps[(i, i)], dps[(i, j)], dps[(j, j)])
                    if d < dmin:
                        dmin = d
                        sol = [(i, j), c, d]
        else:
            dmin = 5e8
            for i in range(len(vecs)):
                for j in range(i + 1, len(vecs)):
                    if (i, j) not in dps:
                        dps[(i, j)] = vecs[i][j]
                        dps[(j, i)] = dps[(i, j)]
                    if (i, i) not in dps:
                        dps[(i, i)] = vecs[i][i]
                    if (j, j) not in dps:
                        dps[(j, j)] = vecs[j][j]
                    c, d = MinNormSolver._min_norm_element_from2(dps[(i, i)], dps[(i, j)], dps[(j, j)])
                    if d < dmin:
                        dmin = d
                        sol = [(i, j), c, d]
            # print("dps", dps)
        return sol, dps

    def _projection2simplex(y):
        m = len(y)
        # print("torch.sort(y)", torch.sort(y)[0])
        sorted_y = torch.flip(torch.sort(y)[0], dims=[0])
        tmpsum = 0.0
        tmax_f = (y.sum() - 1.0) / m
        for i in range(m - 1):
            tmpsum += sorted_y[i]
            tmax = (tmpsum - 1) / (i + 1.0)
            if tmax > sorted_y[i + 1]:
                tmax_f = tmax
                break
        return torch.max(y - tmax_f, torch.zeros(y.shape).cuda())

    def _next_point(cur_val, grad, n):
        proj_grad = grad - (torch.sum(grad) / n)
        tm1 = -1.0 * cur_val[proj_grad < 0] / proj_grad[proj_grad < 0]
        tm2 = (1.0 - cur_val[proj_grad > 0]) / (proj_grad[proj_grad > 0])
        t = 1

        if len(tm1[tm1 > 1e-7]) > 0:
            t = (tm1[tm1 > 1e-7]).min()
        if len(tm2[tm2 > 1e-7]) > 0:
            t = min(t, (tm2[tm2 > 1e-7]).min())

        next_point = proj_grad * t + cur_val
        next_point = MinNormSolver._projection2simplex(next_point)
        return next_point

    def find_min_norm_element(vecs, decomposition):

        # Solution lying at the combination of two points
        dps = {}
        init_sol, dps = MinNormSolver._min_norm_2d(vecs, dps, decomposition)

        n = len(vecs)
        sol_vec = torch.zeros(n).cuda()
        sol_vec[init_sol[0][0]] = init_sol[1]
        sol_vec[init_sol[0][1]] = 1 - init_sol[1]

        if n < 3:
            # This is optimal for n=2, so return the solution
            return sol_vec, init_sol[2]

        iter_count = 0

        grad_mat = torch.zeros((n, n)).cuda()
        for i in range(n):
            for j in range(n):
                grad_mat[i, j] = dps[(i, j)]

        while iter_count < MinNormSolver.MAX_ITER:

            grad_dir = -1.0 * torch.mm(grad_mat, sol_vec.view(-1, 1)).view(-1)
            new_point = MinNormSolver._next_point(sol_vec, grad_dir, n)
            # Re-compute the inner products for line search
            v1v1 = 0.0
            v1v2 = 0.0
            v2v2 = 0.0
            for i in range(n):
                for j in range(n):
                    v1v1 += sol_vec[i] * sol_vec[j] * dps[(i, j)]
                    v1v2 += sol_vec[i] * new_point[j] * dps[(i, j)]
                    v2v2 += new_point[i] * new_point[j] * dps[(i, j)]
            nc, nd = MinNormSolver._min_norm_element_from2(v1v1, v1v2, v2v2)

            new_sol_vec = nc * sol_vec + (1 - nc) * new_point
            change = new_sol_vec - sol_vec
            # print("Change: ", change)
            try:
                if change.pow(2).sum() < MinNormSolver.STOP_CRIT:
                    return sol_vec, nd
            except Exception as e:
                print(e)
                print("Change: ", change)
                # return sol_vec, nd
            sol_vec = new_sol_vec
        return sol_vec, nd


def gradient_normalizers(grads, losses, normalization_type):
    gn = {}
    if normalization_type == "l2":
        for t in range(len(grads)):
            gn[t] = np.sqrt(np.sum([gr.pow(2).sum().data.cpu() for gr in grads[t]]))
    elif normalization_type == "loss":
        for t in range(len(grads)):
            gn[t] = losses[t]
    elif normalization_type == "loss+":
        for t in range(len(grads)):
            gn[t] = losses[t] * np.sqrt(np.sum([gr.pow(2).sum().data.cpu() for gr in grads[t]]))
    elif normalization_type == "none":
        for t in range(len(grads)):
            gn[t] = 1.0
    else:
        print("ERROR: Invalid Normalization Type")
    return gn


class RBF(torch.nn.Module):
    def __init__(self, sigma=None):
        super(RBF, self).__init__()

        self.sigma = sigma

    def forward(self, X, Y, scale):
        XX = X.matmul(X.t())
        XY = X.matmul(Y.t())
        YY = Y.matmul(Y.t())

        dnorm2 = -2 * XY + XX.diag().unsqueeze(1) + YY.diag().unsqueeze(0)

        # Apply the median heuristic (PyTorch does not give true median)
        if self.sigma is None:
            np_dnorm2 = dnorm2.detach().cpu().numpy()
            h = np.median(np_dnorm2) / (2 * np.log(X.size(0) + 1))
            sigma = np.sqrt(h).item()
        else:
            sigma = self.sigma
        sigma = sigma * scale
        gamma = 1.0 / (1e-8 + 2 * sigma ** 2)
        K_XY = (-gamma * dnorm2).exp()

        return K_XY


# Let us initialize a reusable instance right away.
K = RBF()


class MoG(torch.distributions.Distribution):
    def __init__(self, pi, loc, covariance_matrix):
        self.num_components = loc.size(0)
        self.loc = loc
        self.covariance_matrix = covariance_matrix
        self.pi = pi
        self.dists = [
            torch.distributions.MultivariateNormal(mu, covariance_matrix=sigma)
            for mu, sigma in zip(loc, covariance_matrix)
        ]

        super(MoG, self).__init__(torch.Size([]), torch.Size([loc.size(-1)]))

    @property
    def arg_constraints(self):
        return self.dists[0].arg_constraints

    @property
    def support(self):
        return self.dists[0].support

    @property
    def has_rsample(self):
        return False

    def log_prob(self, value):
        res = torch.cat(
            [
                (torch.log(torch.tensor(self.pi[i])) + self.dists[i].log_prob(value)).unsqueeze(-1)
                for i in range(len(self.dists))
            ],
            dim=-1,
        ).logsumexp(dim=-1)
        #     print(res.shape)
        return res

    def enumerate_support(self):
        return self.dists[0].enumerate_support()


class MoG2(MoG):
    def __init__(self, pi, loc, device=None):
        loc = torch.Tensor(loc).to(device)
        cov = torch.Tensor([0.5, 0.5]).diag().unsqueeze(0).repeat(2, 1, 1).to(device)

        super(MoG2, self).__init__(pi, loc, cov)


mog2_1 = MoG2(pi=[0.7, 0.3], loc=[[4.0, -4.0], [0, 0.5]], device=device)
mog2_2 = MoG2(pi=[0.7, 0.3], loc=[[-4.0, 4.0], [0.5, 0.0]], device=device)
mog2_3 = MoG2(pi=[0.7, 0.3], loc=[[-3.0, -3.0], [0.0, 0.0]], device=device)


n = 30
X_init = (5 * torch.randn(n, *mog2_1.event_shape)).to(device)
X_init.data = torch.clamp(X_init.data.clone(), min=-10 + 1e-3, max=10 - 1e-3)


class MTS_SVGD:
    def __init__(self, P, K, optimizer, nornalize="none"):
        self.P = P
        self.K = K
        self.n_targets = len(P)
        self.optim = optimizer
        self.nornalize = nornalize

    def phi(self, X, index, scale=1):
        X = X.detach().requires_grad_(True)

        log_prob = self.P[index].log_prob(X)
        #         print("log_prob", log_prob.shape)
        score_func = autograd.grad(log_prob.sum(), X)[0]  # n_particles x dimension
        #         print("score_func", score_func.shape)

        K_XX = self.K(X, X.detach(), scale=1)
        grad_K = -autograd.grad(K_XX.sum(), X)[0]

        phi1 = K_XX.detach().matmul(score_func) / X.size(0)
        phi2 = grad_K / X.size(0)
        phi = phi1 + phi2
        return phi, score_func

    def step(self, X, scale=1):
        phis = [[] for _ in range(self.n_targets)]
        score_funcs = []
        losses = []
        for i in range(self.n_targets):
            self.optim.zero_grad()
            phi, score_func = self.phi(X, i, scale)
            phis[i].append(Variable(phi.detach().clone(), requires_grad=False))
            score_func = torch.nn.functional.normalize(score_func, dim=0)
            score_funcs.append(Variable(score_func.detach().clone(), requires_grad=False))

        score_funcs = torch.stack(score_funcs, 0)
        #         score_funcs = torch.nn.functional.normalize(score_funcs, dim=0)
        with torch.no_grad():
            Q = torch.zeros((self.n_targets, self.n_targets))
            K_XX = self.K(X, X, scale).detach()

            for i in range(self.n_targets):
                for j in range(i, self.n_targets):

                    Q[i][j] = torch.mul(K_XX, torch.matmul(score_funcs[i], score_funcs[j].T)).sum()
                    Q[j][i] = Q[i][j]

            gn = gradient_normalizers(phis, losses, self.nornalize)
            for i in range(self.n_targets):
                for gr_i in range(len(phis[i])):
                    phis[i][gr_i] = phis[i][gr_i] / gn[i]

            sol, min_norm = MinNormSolver.find_min_norm_element(Q, decomposition=False)
            scales = []
            for i in range(self.n_targets):
                scales.append(float(sol[i]) * phis[i][0])

        X.grad = -torch.stack(scales).sum(dim=0)
        self.optim.step()


class MOO_SVGD:
    def __init__(self, P, K, optimizer, nornalize="none"):
        self.P = P
        self.K = K
        self.n_targets = len(P)
        self.optim = optimizer
        self.nornalize = nornalize

    def phi(self, x, index, scale=1):
        X = x.detach().requires_grad_(True)

        log_prob = self.P[index].log_prob(X)
        score_func = autograd.grad(log_prob.sum(), X)[0]

        K_XX = self.K(X, X.detach(), scale)
        grad_K = -autograd.grad(K_XX.sum(), X)[0]

        phi1 = K_XX.detach().matmul(score_func) / X.size(0)
        phi2 = grad_K / X.size(0)
        phi = phi1 + phi2
        return phi.detach().clone(), score_func.detach().clone()

    def step(self, x, scale=1):
        X = x.detach().requires_grad_(True)
        n_particles = X.shape[0]

        phis = [[] for _ in range(self.n_targets)]
        score_funcs = []
        losses = []
        for i in range(self.n_targets):
            self.optim.zero_grad()
            phi, score_func = self.phi(X, i, scale)
            phis[i].append(Variable(phi.detach().clone(), requires_grad=False))
            #             score_func = torch.nn.functional.normalize(score_func, dim=0)
            score_funcs.append(Variable(score_func.detach().clone(), requires_grad=False))
            self.optim.zero_grad()
        score_funcs = torch.stack(score_funcs, 0)

        K_XX = self.K(X, X.detach(), scale)
        grad_K = -autograd.grad(K_XX.sum(), X)[0]

        with torch.no_grad():

            gn = gradient_normalizers(phis, losses, self.nornalize)
            for i in range(self.n_targets):
                for gr_i in range(len(phis[i])):
                    phis[i][gr_i] = phis[i][gr_i] / gn[i]
            grads = []

            for i in range(n_particles):
                sol, min_norm = MinNormSolver.find_min_norm_element(score_funcs[:, i, :], decomposition=True)
                grads.append(sum([sol[j] * score_funcs[j, i, :] for j in range(self.n_targets)]))

            grads = torch.stack(grads, 0)

            phi1 = K_XX.detach().matmul(grads) / X.size(0)
            phi2 = grad_K / X.size(0)
            phi = phi1 + phi2
        x.grad = -phi
        self.optim.step()


mog2_chart = get_density_charts([mog2_3, mog2_2, mog2_1], d=10.0, step=0.1)

n = 3
X_init = (5 * torch.randn(n, *mog2_1.event_shape)).to(device)
X_init.data = torch.clamp(X_init.data.clone(), min=-10 + 1e-3, max=10 - 1e-3)

MTS_Xs = []
X = X_init.clone()
svgd = MTS_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
for i in tqdm(range(2001), total=2001):
    if i in [0, 250, 500, 1000]:
        MTS_Xs.append(X.detach().clone().cpu().numpy())
    svgd.step(X)

chart1 = mog2_chart + get_particles_chart(MTS_Xs[0])
chart2 = mog2_chart + get_particles_chart(MTS_Xs[1])
chart3 = mog2_chart + get_particles_chart(MTS_Xs[2])
chart4 = mog2_chart + get_particles_chart(MTS_Xs[3])

chart1.title = "Step 0"
chart2.title = "Step 250"
chart3.title = "Step 500"
chart4.title = "Step 1000"

chart = (
    alt.hconcat(chart1, chart2, chart3, chart4, center=True)
    .configure_title(fontSize=20)
    .configure_axis(titleFontSize=16)
)
chart.save("MT-SGD.png")
chart.save("MT-SGD.pdf")
chart

n = 5
X_init = (5 * torch.randn(n, *mog2_1.event_shape)).to(device)
X_init.data = torch.clamp(X_init.data.clone(), min=-10 + 1e-3, max=10 - 1e-3)

MTS_Xs = []
X = X_init.clone()
svgd = MTS_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
for i in tqdm(range(2001), total=2001):
    if i in [0, 250, 500, 1000]:
        MTS_Xs.append(X.detach().clone().cpu().numpy())
    svgd.step(X)

chart1 = mog2_chart + get_particles_chart(MTS_Xs[0])
chart2 = mog2_chart + get_particles_chart(MTS_Xs[1])
chart3 = mog2_chart + get_particles_chart(MTS_Xs[2])
chart4 = mog2_chart + get_particles_chart(MTS_Xs[3])

chart1.title = "Step 0"
chart2.title = "Step 250"
chart3.title = "Step 500"
chart4.title = "Step 1000"

chart = (
    alt.hconcat(chart1, chart2, chart3, chart4, center=True)
    .configure_title(fontSize=20)
    .configure_axis(titleFontSize=16)
)
chart.save("MT-SGD.png")
chart.save("MT-SGD.pdf")
chart

n = 10
X_init = (5 * torch.randn(n, *mog2_1.event_shape)).to(device)
X_init.data = torch.clamp(X_init.data.clone(), min=-10 + 1e-3, max=10 - 1e-3)

MTS_Xs = []
X = X_init.clone()
svgd = MTS_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
for i in tqdm(range(2001), total=2001):
    if i in [0, 250, 500, 1000]:
        MTS_Xs.append(X.detach().clone().cpu().numpy())
    svgd.step(X)

chart1 = mog2_chart + get_particles_chart(MTS_Xs[0])
chart2 = mog2_chart + get_particles_chart(MTS_Xs[1])
chart3 = mog2_chart + get_particles_chart(MTS_Xs[2])
chart4 = mog2_chart + get_particles_chart(MTS_Xs[3])

chart1.title = "Step 0"
chart2.title = "Step 250"
chart3.title = "Step 500"
chart4.title = "Step 1000"

chart = (
    alt.hconcat(chart1, chart2, chart3, chart4, center=True)
    .configure_title(fontSize=20)
    .configure_axis(titleFontSize=16)
)
chart.save("MT-SGD.png")
chart.save("MT-SGD.pdf")
chart

n = 50
X_init = (5 * torch.randn(n, *mog2_1.event_shape)).to(device)
X_init.data = torch.clamp(X_init.data.clone(), min=-10 + 1e-3, max=10 - 1e-3)

MTS_Xs = []
X = X_init.clone()
svgd = MTS_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
for i in tqdm(range(2001), total=2001):
    if i in [0, 250, 500, 1000]:
        MTS_Xs.append(X.detach().clone().cpu().numpy())
    svgd.step(X)

chart1 = mog2_chart + get_particles_chart(MTS_Xs[0])
chart2 = mog2_chart + get_particles_chart(MTS_Xs[1])
chart3 = mog2_chart + get_particles_chart(MTS_Xs[2])
chart4 = mog2_chart + get_particles_chart(MTS_Xs[3])

chart1.title = "Step 0"
chart2.title = "Step 250"
chart3.title = "Step 500"
chart4.title = "Step 1000"

chart = (
    alt.hconcat(chart1, chart2, chart3, chart4, center=True)
    .configure_title(fontSize=20)
    .configure_axis(titleFontSize=16)
)
chart.save("MT-SGD.png")
chart.save("MT-SGD.pdf")
chart

MOO_Xs = []
X = X_init.clone()
svgd = MOO_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
for i in tqdm(range(2001), total=2001):
    if i in [0, 500, 1000, 2000]:
        MOO_Xs.append(X.detach().clone().cpu().numpy())
    svgd.step(X)


chart1 = mog2_chart + get_particles_chart(MOO_Xs[0])
chart2 = mog2_chart + get_particles_chart(MOO_Xs[1])
chart3 = mog2_chart + get_particles_chart(MOO_Xs[2])
chart4 = mog2_chart + get_particles_chart(MOO_Xs[3])

chart1.title = "Step 0"
chart2.title = "Step 500"
chart3.title = "Step 1000"
chart4.title = "Step 2000"

chart = (
    alt.hconcat(chart1, chart2, chart3, chart4, center=True)
    .configure_title(fontSize=20)
    .configure_axis(titleFontSize=16)
)
chart.save("MOO-SVGD.png")
chart.save("MOO-SVGD.pdf")
chart

for _ in range(5):
    n = 5
    X_init = (5 * torch.randn(n, *mog2_1.event_shape)).to(device)
    X_init.data = torch.clamp(X_init.data.clone(), min=-10 + 1e-3, max=10 - 1e-3)

    X = X_init.clone()
    svgd = MTS_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
    for i in tqdm(range(1001), total=1001):
        svgd.step(X)

    X = X_init.clone()
    svgd = MOO_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
    for i in tqdm(range(1001), total=1001):
        svgd.step(X)

for _ in range(5):

    n = 10
    X_init = (5 * torch.randn(n, *mog2_1.event_shape)).to(device)
    X_init.data = torch.clamp(X_init.data.clone(), min=-10 + 1e-3, max=10 - 1e-3)

    X = X_init.clone()
    svgd = MTS_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
    for i in tqdm(range(1001), total=1001):
        svgd.step(X)

    X = X_init.clone()
    svgd = MOO_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
    for i in tqdm(range(1001), total=1001):
        svgd.step(X)

for _ in range(5):

    n = 25
    X_init = (5 * torch.randn(n, *mog2_1.event_shape)).to(device)
    X_init.data = torch.clamp(X_init.data.clone(), min=-10 + 1e-3, max=10 - 1e-3)

    X = X_init.clone()
    svgd = MTS_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
    for i in tqdm(range(1001), total=1001):
        svgd.step(X)

    X = X_init.clone()
    svgd = MOO_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
    for i in tqdm(range(1001), total=1001):
        svgd.step(X)

for _ in range(5):
    n = 50
    X_init = (5 * torch.randn(n, *mog2_1.event_shape)).to(device)
    X_init.data = torch.clamp(X_init.data.clone(), min=-10 + 1e-3, max=10 - 1e-3)

    X = X_init.clone()
    svgd = MTS_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
    for i in tqdm(range(1001), total=1001):
        svgd.step(X)

    X = X_init.clone()
    svgd = MOO_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
    for i in tqdm(range(1001), total=1001):
        svgd.step(X)

for _ in range(5):
    n = 100
    X_init = (5 * torch.randn(n, *mog2_1.event_shape)).to(device)
    X_init.data = torch.clamp(X_init.data.clone(), min=-10 + 1e-3, max=10 - 1e-3)

    X = X_init.clone()
    svgd = MTS_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
    for i in tqdm(range(1001), total=1001):
        svgd.step(X)

    X = X_init.clone()
    svgd = MOO_SVGD([mog2_3, mog2_2, mog2_1], K, optim.Adam([X], lr=3e-2), nornalize="none")
    for i in tqdm(range(1001), total=1001):
        svgd.step(X)