import torch
import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self, num_channels, num_groups=32):
        super(Model, self).__init__()
        self.gn = nn.GroupNorm(num_groups, num_channels)

    def forward(self, x, y):
        eps = 1e-8
        p = torch.abs(x) + eps
        q = torch.abs(y) + eps
        m = 0.5 * (p + q)

        kl_p = p * (torch.log(p) - torch.log(m))
        kl_q = q * (torch.log(q) - torch.log(m))

        js_elem = 0.5 * (kl_p + kl_q)
        out = self.gn(js_elem)
        return out.mean()


batch_size = 16
input_dim = 1024


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


def get_init_inputs():
    return [input_dim]