support fftshift and ifftshift in pytorch

[Mon-ius]

Currently working on the IBM Power9 series clusters.

Cannot run such code:

x = torch.tensor(1, dtype=torch.complex64, device = 'cuda')

Error shows below:

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-93-353b552465e0> in <module>
----> 1 x = torch.tensor(1, dtype=torch.complex64, device = 'cuda')

RuntimeError: Could not run 'aten::empty.memory_format' with arguments from the 'ComplexCPUTensorId' backend. 'aten::empty.memory_format' is only available for these backends: [CUDATensorId, SparseCPUTensorId, VariableTensorId, CPUTensorId, MkldnnCPUTensorId, SparseCUDATensorId].

cc @ezyang @anjali411 @dylanbespalko

[ngimel]

What pytorch version are you using? The snippet above works on master.

[Mon-ius]

1.4.0

  Operating System: Red Hat Enterprise Linux
       CPE OS Name: cpe:/o:redhat:enterprise_linux:7.6:GA:server
            Kernel: Linux 4.14.0-115.7.1.el7a.ppc64le
      Architecture: ppc64-le

[ngimel]

Please update to the latest version

[Mon-ius]

I tried the latest version. It does support the complex type. But I didn't see there is some function like np.fft.fftshift or np.fft.ifftshift if it has already been supported or I have to write one to do that. I don't want to switch to numpy and pytorch a lot during training.

[zasdfgbnm]

fft support is tracked at #33152

[ArthDh]

@Mon-ius I couldn't find fftshift/ifftshift on master but this should work fine:

source: locuslab/pytorch_fft#12

def roll_n(X, axis, n):
    f_idx = tuple(slice(None, None, None) if i != axis else slice(0, n, None)
                  for i in range(X.dim()))
    b_idx = tuple(slice(None, None, None) if i != axis else slice(n, None, None)
                  for i in range(X.dim()))
    front = X[f_idx]
    back = X[b_idx]
    return torch.cat([back, front], axis)

def fftshift(X):
    # batch*channel*...*2
    real, imag = X.chunk(chunks=2, dim=-1)
    real, imag = real.squeeze(dim=-1), imag.squeeze(dim=-1)

    for dim in range(2, len(real.size())):
        real = roll_n(real, axis=dim, n=int(np.ceil(real.size(dim) / 2)))
        imag = roll_n(imag, axis=dim, n=int(np.ceil(imag.size(dim) / 2)))

    real, imag = real.unsqueeze(dim=-1), imag.unsqueeze(dim=-1)
    X = torch.cat((real,imag),dim=-1)
    return X

def ifftshift(X):
    # batch*channel*...*2
    real, imag = X.chunk(chunks=2, dim=-1)
    real, imag = real.squeeze(dim=-1), imag.squeeze(dim=-1)

    for dim in range(len(real.size()) - 1, 1, -1):
        real = roll_n(real, axis=dim, n=int(np.floor(real.size(dim) / 2)))
        imag = roll_n(imag, axis=dim, n=int(np.floor(imag.size(dim) / 2)))

    real, imag = real.unsqueeze(dim=-1), imag.unsqueeze(dim=-1)
    X = torch.cat((real, imag), dim=-1)

    return X

[ngimel]

Thanks @ArthDh, I'd like to point out that torch has torch.roll function that provides similar funcitonality to roll_n with likely better performance.

[ngimel]

Closing in favor of #42175