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


class Model(nn.Module):
    def __init__(self, gamma=0.5):
        super(Model, self).__init__()
        self.gamma = gamma

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

        sum_p_pow = torch.sum(p_prob.pow(1.0 + self.gamma), dim=1)
        sum_pq_pow = torch.sum(p_prob * q_prob.pow(self.gamma), dim=1)
        sum_q_pow = torch.sum(q_prob.pow(1.0 + self.gamma), dim=1)

        term1 = torch.log(sum_p_pow) / (self.gamma * (1.0 + self.gamma))
        term2 = torch.log(sum_pq_pow) / self.gamma
        term3 = torch.log(sum_q_pow) / (1.0 + self.gamma)

        loss = term1 - term2 + term3
        return loss.mean()


batch_size = 32
num_classes = 1000


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


def get_init_inputs():
    return [0.5]