Implement n-dimensional hermitian FFTs #63890

peterbell10 · 2021-08-24T21:20:17Z

Closes #59127

cc @mruberry @peterbell10 @walterddr

facebook-github-bot · 2021-08-24T21:20:22Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/63890
📄 Preview docs built from this PR
📄 Preview C++ docs built from this PR
🔧 Opt-in to CIFlow to control what jobs run on your PRs

💊 CI failures summary and remediations

As of commit 7a01f5b (more details on the Dr. CI page):

1/1 failures introduced in this PR

1 failure not recognized by patterns:

Job	Step	Action
^{Lint / clang-format}	^{Run clang-format}	🔁 rerun

This comment was automatically generated by Dr. CI (expand for details).

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

mruberry · 2021-08-26T03:50:47Z

This is really exciting but it looks like the test failures are real

codecov · 2021-08-26T16:16:59Z

Codecov Report

Merging #63890 (16cda64) into master (b0782f0) will decrease coverage by 0.26%.
The diff coverage is 100.00%.

❗ Current head 16cda64 differs from pull request most recent head f10778f. Consider uploading reports for the commit f10778f to get more accurate results

@@            Coverage Diff             @@
##           master   #63890      +/-   ##
==========================================
- Coverage   67.03%   66.76%   -0.27%     
==========================================
  Files         695      695              
  Lines       90643    90749     +106     
==========================================
- Hits        60760    60588     -172     
- Misses      29883    30161     +278

peterbell10 · 2021-08-26T16:38:04Z

@mruberry PTAL, I've fixed the tests.

mruberry · 2021-08-31T06:08:04Z

aten/src/ATen/native/native_functions.yaml

@@ -10135,6 +10135,26 @@
  python_module: fft
  variants: function

+- func: fft_hfft2(Tensor self, int[1]? s=None, int[1] dim=[-2,-1], str? norm=None) -> Tensor


Nice following of SciPy's hfft2 signature

mruberry · 2021-08-31T06:15:34Z

aten/src/ATen/native/native_functions.yaml

+  python_module: fft
+  variants: function
+
+- func: fft_ihfft2(Tensor self, int[1]? s=None, int[1] dim=[-2,-1], str? norm=None) -> Tensor


ihfft2 signature looks great, too

mruberry · 2021-08-31T06:18:21Z

aten/src/ATen/native/native_functions.yaml

@@ -10167,6 +10187,26 @@
  python_module: fft
  variants: function

+- func: fft_hfftn(Tensor self, int[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor


hfftn signature looks great

mruberry · 2021-08-31T06:19:01Z

aten/src/ATen/native/native_functions.yaml

+  python_module: fft
+  variants: function
+
+- func: fft_ihfftn(Tensor self, int[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor


mruberry · 2021-08-31T06:19:50Z

test/test_spectral_ops.py

@@ -25,6 +25,20 @@
 if TEST_LIBROSA:
    import librosa

+has_scipy_fft = False


mruberry · 2021-08-31T06:20:20Z

test/test_spectral_ops.py

+except ModuleNotFoundError:
+    pass
+
+LooseVersion = distutils.version.LooseVersion


Nice versioning

mruberry · 2021-08-31T06:23:36Z

torch/testing/_internal/common_methods_invocations.py

@@ -6272,6 +6278,18 @@ def gradcheck_wrapper_triangular_input(op, input, *args, upper=False, **kwargs):
                     dtypes=all_types_and_complex_and(torch.bool),
                     default_test_dtypes=floating_and_complex_types(),
                     check_batched_gradgrad=False),
+    SpectralFuncInfo('fft.hfftn',


For completeness we should add hfft2 and ihfft2, too (this can be done in a follow-up PR)

mruberry · 2021-08-31T06:24:38Z

torch/fft/__init__.py

+hfftn = _add_docstr(_fft.fft_hfftn, r"""
+hfftn(input, s=None, dim=None, norm=None, *, out=None) -> Tensor
+
+Computes the n-dimensional discrete Fourier transform of a Herimitian symmetric


Great first sentence

mruberry · 2021-08-31T06:24:45Z

torch/fft/__init__.py

+Computes the n-dimensional discrete Fourier transform of a Herimitian symmetric
+:attr:`input` signal.
+
+:attr:`input` is interpreted as a one-sided Hermitian signal in the time


mruberry · 2021-08-31T06:25:48Z

torch/fft/__init__.py

+    a real signal in the time-domain and gives Hermitian symmetry in the
+    frequency-domain. The Hermitian FFT is the opposite; Hermitian symmetric in
+    the time-domain and real-valued in the frequency-domain. For this reason,
+    special care needs to be taken with the shape argument :attr:`s`, in the


This last sentence is a little vague but good enough for now

mruberry · 2021-08-31T06:26:29Z

torch/fft/__init__.py

+    the original data, as given by :attr:`s`. This is because each input shape
+    could correspond to either an odd or even length signal. By default, the
+    signal is assumed to be even length and odd signals will not round-trip
+    properly. So, it is recommended to always pass the signal shape :attr:`s`.


"So," can be removed and the last sentence can start with "It is..."

torch/fft/__init__.py

mruberry · 2021-08-31T06:30:30Z

torch/fft/__init__.py

+    dim (Tuple[int], optional): Dimensions to be transformed.
+        The last dimension must be the half-Hermitian compressed dimension.
+        Default: all dimensions, or the last ``len(s)`` dimensions if :attr:`s` is given.
+    norm (str, optional): Normalization mode. For the forward transform


Great norm description

mruberry · 2021-08-31T06:33:32Z

torch/fft/__init__.py

+
+Example:
+
+    >>> t = torch.rand(10, 10)


This example is cool but doesn't actually call ihfftn anywhere

mruberry · 2021-08-31T06:34:10Z

torch/fft/__init__.py

+ihfft2 = _add_docstr(_fft.fft_ihfft2, r"""
+ihfft2(input, s=None, dim=(-2, -1), norm=None, *, out=None) -> Tensor
+
+Computes the N-dimensional inverse discrete Fourier transform of real


Great description

mruberry · 2021-08-31T06:34:24Z

torch/fft/__init__.py

+    torch.Size([10, 6])
+
+    Compared against the full output from :func:`~torch.fft.ifft2`, the
+    hermitian time-space signal takes up only half the space.


mruberry · 2021-08-31T06:34:42Z

torch/fft/__init__.py

+
+    >>> two_ffts = torch.fft.ifft(torch.fft.ihfft(t, dim=1), dim=0)
+    >>> torch.allclose(t, two_ffts)
+


mruberry

Hey @peterbell10! This looks great overall. I made a few doc comments for your review. Nothing major but I think one example is still a placeholder. Just ping me when this is ready to go!

peterbell10 · 2021-08-31T22:52:39Z

@mruberry PTAL

mruberry · 2021-09-06T06:05:29Z

aten/src/ATen/native/SpectralOps.cpp

-  desc.shape.back() = last_dim_size / 2 + 1;
+  auto ld = last_dim_size / 2 + 1;
+  desc.shape.back() = ld;
+  TORCH_CHECK(ld >= 1, "Invalid number of data points (", last_dim_size, ") specified");


fname should probably sneak in here, too

mruberry

Awesome! Thanks @peterbell10!

Did you see that JAX added SciPy's dct and dctn? Might be time to revisit #49016.

Looking at uses within Facebook, I can there are some uses of dct and dctn internally, and these ops seem reasonably popular on Github, too.

facebook-github-bot · 2021-09-06T06:21:34Z

@mruberry has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

ngimel · 2021-09-06T18:59:26Z

test/test_spectral_ops.py

+        # input_ndim, dim
+        transform_desc = [
+            *product(range(2, 5), (None, (0,), (0, -1))),
+            *product(range(2, 5), (None,)),


nit: why is this line needed? (2,None) etc that it generates is already generated by the previous line.

ngimel · 2021-09-06T18:59:59Z

test/test_spectral_ops.py

+            (3, (0, -1)),
+            (3, (1,)),
+            (1, (0,)),
+            (4, None),


ngimel · 2021-09-06T19:00:05Z

test/test_spectral_ops.py

+            *product(range(2, 5), (None,)),
+            (6, None),
+            (5, (1, 3, 4)),
+            (3, (0, -1)),


ngimel · 2021-09-06T19:01:32Z

torch/fft/__init__.py

+hfft2 = _add_docstr(_fft.fft_hfft2, r"""
+hfft2(input, s=None, dim=(-2, -1), norm=None, *, out=None) -> Tensor
+
+Computes the 2-dimensional discrete Fourier transform of a Herimitian symmetric


typo: Herimitian -> Hermitian

ngimel · 2021-09-06T19:12:34Z

torch/fft/__init__.py

+Computes the 2-dimensional discrete Fourier transform of a Herimitian symmetric
+:attr:`input` signal. Equivalent to :func:`~torch.fft.hfftn` but only
+transforms the last two dimensions by default.
+


I think the note from hfftn that input is interpreted as a one-sided Hermitian signal in the time domain should also be repeated here, otherwise it's unclear what if input itself is required to have some symmetric properties.

mruberry · 2021-09-13T01:44:05Z

@peterbell10 just ping me when @ngimel's comments are addressed (thanks for taking a look, @ngimel!)

Remove redundant test parameters and doc tweaks.

pytorch-probot · 2021-09-20T10:27:58Z

CI Flow Status

⚛️ CI Flow

Ruleset - Version: v1
Ruleset - File: https://github.com/peterbell10/pytorch/blob/7a01f5b06c92854c290d328833f5ce905df09ab6/.github/generated-ciflow-ruleset.json
PR ciflow labels: ciflow/default

Workflows	Labels (bold enabled)	Status
Triggered Workflows
linux-bionic-py3.6-clang9	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`, `ciflow/noarch`, `ciflow/xla`	✅ triggered
linux-bionic-py3.8-gcc9-coverage	`ciflow/all`, `ciflow/coverage`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`	✅ triggered
linux-xenial-cuda11.3-py3.6-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/default`, `ciflow/linux`	✅ triggered
linux-xenial-py3.6-gcc5.4	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`	✅ triggered
linux-xenial-py3.6-gcc7-bazel-test	`ciflow/all`, `ciflow/bazel`, `ciflow/cpu`, `ciflow/default`, `ciflow/linux`	✅ triggered
win-vs2019-cpu-py3	`ciflow/all`, `ciflow/cpu`, `ciflow/default`, `ciflow/win`	✅ triggered
win-vs2019-cuda11.3-py3	`ciflow/all`, `ciflow/cuda`, `ciflow/default`, `ciflow/win`	✅ triggered
Skipped Workflows
libtorch-linux-xenial-cuda10.2-py3.6-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/libtorch`, `ciflow/linux`	🚫 skipped
libtorch-linux-xenial-cuda11.3-py3.6-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/libtorch`, `ciflow/linux`	🚫 skipped
linux-bionic-cuda10.2-py3.9-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/linux`, `ciflow/slow`	🚫 skipped
linux-xenial-cuda10.2-py3.6-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/linux`, `ciflow/slow`	🚫 skipped
parallelnative-linux-xenial-py3.6-gcc5.4	`ciflow/all`, `ciflow/cpu`, `ciflow/linux`	🚫 skipped
periodic-libtorch-linux-xenial-cuda11.1-py3.6-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/libtorch`, `ciflow/linux`, `ciflow/scheduled`	🚫 skipped
periodic-linux-xenial-cuda11.1-py3.6-gcc7	`ciflow/all`, `ciflow/cuda`, `ciflow/linux`, `ciflow/scheduled`	🚫 skipped
periodic-win-vs2019-cuda11.1-py3	`ciflow/all`, `ciflow/cuda`, `ciflow/scheduled`, `ciflow/win`	🚫 skipped
puretorch-linux-xenial-py3.6-gcc5.4	`ciflow/all`, `ciflow/cpu`, `ciflow/linux`	🚫 skipped
win-vs2019-cuda10.2-py3	`ciflow/all`, `ciflow/cuda`, `ciflow/win`	🚫 skipped

You can add a comment to the PR and tag @pytorchbot with the following commands:

# ciflow rerun, "ciflow/default" will always be added automatically
@pytorchbot ciflow rerun

# ciflow rerun with additional labels "-l <ciflow/label_name>", which is equivalent to adding these labels manually and trigger the rerun
@pytorchbot ciflow rerun -l ciflow/scheduled -l ciflow/slow

For more information, please take a look at the CI Flow Wiki.

peterbell10 · 2021-09-22T23:07:11Z

@mruberry PTAL.

facebook-github-bot · 2021-09-22T23:47:16Z

@mruberry has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot · 2021-09-23T07:48:44Z

@mruberry has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

peterbell10 · 2021-09-30T14:41:59Z

@mruberry any update on this?

mruberry · 2021-09-30T22:58:28Z

It's just been hitting some internal unrelated land failures. I'm jedi landing it now to bypass them

peterbell10 added open source module: fft labels Aug 24, 2021

peterbell10 requested a review from ezyang as a code owner August 24, 2021 21:20

facebook-github-bot added the cla signed label Aug 24, 2021

peterbell10 force-pushed the hermitian-fftn branch from 5d91bd4 to eeca95d Compare August 24, 2021 21:36

peterbell10 requested a review from mruberry August 24, 2021 21:41

heitorschueroff added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Aug 24, 2021

peterbell10 force-pushed the hermitian-fftn branch from d0cee7d to b701a60 Compare August 26, 2021 01:27

peterbell10 force-pushed the hermitian-fftn branch 2 times, most recently from c04a026 to 16cda64 Compare August 26, 2021 11:48