import torch


def complex_mult(data1, data2, dim=-1):
    """
    Element-wise complex matrix multiplication X^T Y

    :param data1 -> torch.Tensor:
    :param data2 -> torch.Tensor:
    :param dim -> int: dimension that represents the complex values
    """

    assert data1.size(dim) == 2
    assert data2.size(dim) == 2
    re1, im1 = torch.unbind(data1, dim=dim)
    re2, im2 = torch.unbind(data2, dim=dim)

    return torch.stack([re1 * re2 - im1 * im2, im1 *re2 + re1 * im2], dim=dim)