RFC: specifying the shape of the transformed output in n-dimensional FFT APIs

[kgryte]

This RFC proposes that n-dimensional FFT APIs should support specifying a placeholder value when specifying the shape of the transformed output.

Prior Art

In PyTorch, when specifying the shape of the transformed output, one can specify a placeholder value of -1 to indicate that a dimension should not be padded.

>>>  a = torch.tensor([[[0, 0, 0],
...         [0, 0, 0],
...         [0, 0, 0]],
... 
...        [[1, 1, 1],
...         [1, 1, 1],
...         [1, 1, 1]],
... 
...        [[2, 2, 2],
...         [2, 2, 2],
...         [2, 2, 2]]])
>>> y = torch.fft.fftn(a, s=(a.shape[1],3))
>>> y
tensor([[[ 0.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j]],

        [[ 9.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j]],

        [[18.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j]]])
>>> y = torch.fft.fftn(a, s=(-1,3))
>>> y
tensor([[[ 0.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j]],

        [[ 9.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j]],

        [[18.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j],
         [ 0.+0.j,  0.+0.j,  0.+0.j]]])

This is especially convenient when not wanting to access the array shape (as might be the case for accelerator libraries).

In contrast, NumPy does not support a placeholder, requiring that, in order for a dimension to not be padded, one must specify the existing dimension size.

>>> np.fft.fftn(a, s=(3, 3))
array([[[ 0.+0.j,  0.+0.j,  0.+0.j],
        [ 0.+0.j,  0.+0.j,  0.+0.j],
        [ 0.+0.j,  0.+0.j,  0.+0.j]],

       [[ 9.+0.j,  0.+0.j,  0.+0.j],
        [ 0.+0.j,  0.+0.j,  0.+0.j],
        [ 0.+0.j,  0.+0.j,  0.+0.j]],

       [[18.+0.j,  0.+0.j,  0.+0.j],
        [ 0.+0.j,  0.+0.j,  0.+0.j],
        [ 0.+0.j,  0.+0.j,  0.+0.j]]])
>>> np.fft.fftn(a, s=(-1, 3))
Traceback (most recent call last):
  File "/.../numpy/fft/_pocketfft.py", line 80, in _get_forward_norm
    raise ValueError(f"Invalid number of FFT data points ({n}) specified.")
ValueError: Invalid number of FFT data points (-1) specified.

Proposal

This RFC proposes that the proposed FFT specification adopt PyTorch's API and support -1 as a valid value to indicate that a dimension should not be padded (or truncated).

Questions

• Does this RFC present overriding backward compatibility concerns?
• Is there an alternative approach which would be better?

[rgommers]

As far as I can tell, this is a minor convenience feature and it doesn't matter much either way. It will require a bit of extra checking for parsing s, but given that in the majority of usage it is expected for s to be None the extra overhead doesn't matter I'd say.

@peterbell10 do you have an opinion on this, since you implemented it in PyTorch? Are you happy for it to be added to numpy.fft and scipy.fft?

[peterbell10]

This is already implemented in scipy.fft and dates back to scipy.fftpack. numpy.fft does support having None values inside the tuple as a placeholder, but I don't think this is documented anywhere. It's just a side effect of NumPy's fftn calling fft in a loop, which in turn supports n=None.

I will say that supporting s=(None,) in PyTorch wouldn't be possible currently, so having -1 be the standard placeholder works better there.

[leofang]

I will say that supporting s=(None,) in PyTorch wouldn't be possible currently

Thanks for sharing, @peterbell10. What would be the technical issue here?

[peterbell10]

My wording was a bit too strong, it's possible but would require either some significant work to implement List[Optional[T]] in PyTorch's C++ code-generation, or require extra python wrapper code that would be a divergence from the C++ torch api.