fft.py

import torch
from torch import Tensor
from torch.testing import assert_close


def fft(tensor: Tensor, pi: Tensor = torch.acos(torch.tensor(-1.))) -> Tensor:
    n = tensor.size()[-1]
    index = torch.arange(n, dtype=torch.float32)
    fourier = torch.exp(index[:, None] * index[None, :] * pi * -2j / n)

    return tensor.to(dtype=fourier.dtype) @ fourier


a = torch.randn((1, 4))
assert_close(fft(a), torch.fft.fft(a, dim=-1))