import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self, p_re, p_im):
        super(Model, self).__init__()
        self.p_c = torch.complex(torch.tensor(p_re), torch.tensor(p_im))

    def forward(self, z_f):
        z_c = torch.complex(z_f[..., 0], z_f[..., 1])

        y_c = torch.exp(z_c)
        w_c = torch.log(y_c)
        out_c = torch.pow(w_c, self.p_c)

        return torch.stack([out_c.real, out_c.imag], dim=-1)


batch_size = 1024
dim = 2


def get_inputs():
    z_f = torch.randn(batch_size, dim)
    return [z_f]


def get_init_inputs():
    return [0.5, 0.5]