使用numpy模拟torch.fft.fft拯救paddle

import numpy as np
import torch
import paddle
def paddle_fft(x,dim=-1):
    if dim==-1:
        return  paddle.to_tensor(np.fft.fft(x.numpy()))
    else:
        shape= [i for i in range(len(x.shape))]
        shape[dim],shape[-1]=shape[-1],shape[dim]

        x=np.transpose(np.fft.fft(np.transpose(x.numpy(), shape)),shape)
        return paddle.to_tensor(x)





if __name__ == '__main__':
    data=paddle.to_tensor(np.array([[[1, 4, 3], [1, 2, 3]], [[1, 2, 3], [1, 2, 3]]]))

    paddle_f_d=paddle_fft(paddle_fft(data,-1),-2)
    torch_f_d =paddle_fft(torch.fft.fft(torch.Tensor(data.numpy()),dim=-1),-2)
    print(paddle_f_d.numpy())
    print(torch_f_d.numpy())