import torch


def imag_exp(a, dim=0):
    """Imaginary exponential, exp(ia), returns real/imag separate in dim.

    Args:
        a (tensor): A tensor where dimension dim is the complex dimension.
        dim (int, default=0): An integer indicating the complex dimension of
            the output.

    Returns:
        tensor: c = exp(i*a), where i is sqrt(-1).
    """
    c = torch.stack((torch.cos(a), torch.sin(a)), dim)

    return c