import torch
import torch.nn as nn


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

    def forward(self, a, b, c):
        a_c = torch.complex(a[..., 0], a[..., 1])
        b_c = torch.complex(b[..., 0], b[..., 1])
        c_c = torch.complex(c[..., 0], c[..., 1])

        out_c = (torch.conj(a_c) * b_c) / c_c

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


batch_size = 1024


def get_inputs():
    a = torch.randn(batch_size, 2)
    b = torch.randn(batch_size, 2)
    c = torch.randn(batch_size, 2) + 1.0  # Bias away from zero
    return [a, b, c]


def get_init_inputs():
    return []