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:
        eps = 1e-8

        m = 0.5 * (x + self.target)

        kl_x_m = torch.sum(x * (torch.log(x + eps) - torch.log(m + eps)), dim=-1)

        kl_t_m = torch.sum(self.target * (torch.log(self.target + eps) - torch.log(m + eps)), dim=-1)

        js_div = 0.5 * (kl_x_m + kl_t_m)

        return js_div * torch.sigmoid(js_div)


batch_size = 128
input_dim = 1024


def get_inputs():
    # Use softmax to ensure inputs are valid probability distributions (sum to 1, positive)
    x = torch.softmax(torch.randn(batch_size, input_dim), dim=-1)
    return [x]


def get_init_inputs():
    target = torch.softmax(torch.randn(input_dim), dim=-1)
    return [target]