import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, target):
        super(Model, self).__init__()
        self.target = nn.Parameter(target)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_prob = torch.softmax(x, dim=-1)
        t_prob = torch.softmax(self.target, dim=-1)

        min_vals = torch.min(x_prob, t_prob)
        m = torch.sum(min_vals, dim=-1)

        sim_j = (2.0 * m + 1.0) / 3.0

        prefix = torch.sum(min_vals[:, :4], dim=-1)

        return sim_j + 0.1 * prefix * (1.0 - sim_j)


batch_size = 128
input_dim = 1024


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


def get_init_inputs():
    target = torch.randn(input_dim)
    return [target]