import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, dim):
        super(Model, self).__init__()
        self.ln = nn.LayerNorm(dim)

    def forward(self, x, y):
        eps = 1e-8

        x_safe = torch.abs(x) + eps
        y_safe = torch.abs(y) + eps

        kl_elem = y_safe * (torch.log(y_safe) - torch.log(x_safe))
        out = self.ln(kl_elem)
        return out.mean()


batch_size = 16
input_dim = 1024


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


def get_init_inputs():
    return [input_dim]