from torch.autograd import grad

def data(n=10000, d=2, eps=1):
    x = torch.randn(n, d) * eps
    y = x + torch.randn(n, d) * eps
    z = y + torch.randn(n, d)
    return torch.cat((x, z), dim=1), y.sum(dim=1, keepdim=True)

dummy_psi = torch.nn.Parameter(torch.ones(4, 1))
dummy_w = torch.nn.Parameter(torch.Tensor([1.0]))

opt = torch.optim.SGD([phi], lr=1e-3)
mse = torch.nn.MSELoss(reduction="none")

pairs = data(eps=0.1)

for i in range(50000):
    error = 0
    penalty_value = 0
    for x, y in [pairs]:
        error_e = mse(x @ dummy_psi * dummy_w, y)
        mean_error = error_e.mean()
        g = grad(mean_error, dummy_w)[0]
        penalty_value += g.sum()
        error += mean_error
    opt.zero_grad()
    (1e-3 * error + penalty_value).backward()
    opt.step()
    if i % 1000 == 0:
        print(dummy_psi)