import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.epsilon = 0.1
        self.num_iters = 5

    def forward(self, x: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        x_sq = torch.sum(x ** 2, dim=1, keepdim=True)
        t_sq = torch.sum(target ** 2, dim=1, keepdim=True)
        dist_sq = x_sq + t_sq.t() - 2 * torch.matmul(x, target.t())

        C = dist_sq
        M = -C / self.epsilon

        f = torch.zeros(x.size(0), 1, device=x.device)
        g = torch.zeros(1, x.size(0), device=x.device)

        for _ in range(self.num_iters):
            f = -torch.logsumexp(M + g, dim=1, keepdim=True)
            g = -torch.logsumexp(M + f, dim=0, keepdim=True)

        log_P = f + M + g
        P = torch.exp(log_P)

        loss = torch.sum(P * C) / x.size(0)
        return loss


batch_size = 128
input_dim = 1024


def get_inputs():
    x = torch.randn(batch_size, input_dim)
    target = torch.randn(batch_size, input_dim)
    return [x, target]


def get_init_inputs():
    return []