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


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

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

        # sum(P^alpha * Q^(1-alpha))
        sum_term = torch.sum((p_prob ** self.alpha) * (q_prob ** (1.0 - self.alpha)), dim=1)

        # 1/(alpha-1) * log(sum)
        loss = torch.log(sum_term) / (self.alpha - 1.0)

        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]