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


class Model(nn.Module):
    def __init__(self, eps=1e-8):
        super(Model, self).__init__()
        self.eps = eps

    def forward(self, p, target):
        p_prob = F.softmax(p, dim=1)
        q_prob = F.softmax(target, dim=1)

        f_p = p_prob * torch.log(p_prob + self.eps)

        f_q = q_prob * torch.log(q_prob + self.eps)

        df_q = 1.0 + torch.log(q_prob + self.eps)

        term3 = df_q * (p_prob - q_prob)

        divergence = f_p - f_q - term3

        loss = torch.sum(divergence, dim=1)
        return loss.mean()


batch_size = 32
num_classes = 1000


def get_inputs():
    p = torch.randn(batch_size, num_classes, requires_grad=True)
    target = torch.randn(batch_size, num_classes)
    return [p, target]


def get_init_inputs():
    return []