Migrate inverse short-time Fourier transform from torchaudio

[mthrok]

As suggested in #3775, this PR migrates Python code to perform ISTFT from torchaudio to PyTorch.

• Migrate function.
• Migrate tests.
  • round trip precision test
  • Batch test
  • Linearity test
  • Sine test
  • Argument validation test
  • Fundamental case test (zeros and ones)

[dr-ci]

💊 CircleCI build failures summary and remediations

As of commit 45a146f (more details on the Dr. CI page):

---

✅ None of the build failures appear to be your fault 💚

---

• 1/1 broken upstream at merge base b55dee9 from Apr 06 until Apr 07 (9 commits; 2e8f954 - 444073e)

  Please rebase on the viable/strict branch (expand for instructions)

  If your commit is newer than viable/strict, you can try basing on an older, stable commit:

  git fetch https://github.com/pytorch/pytorch viable/strict
  git rebase --onto FETCH_HEAD $(git merge-base origin/master HEAD)
  

  If your commit is older than viable/strict:

  git fetch https://github.com/pytorch/pytorch viable/strict
  git rebase FETCH_HEAD
  

  Check out the recency history of this "viable master" tracking branch.

---

🚧 1 upstream failure:

These were probably caused by upstream breakages:

• pytorch_bazel_build from Apr 06 until Apr 07 (9 commits; 2e8f954 - 444073e)

---

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 on the GitHub issue tracker.

See how this bot performed.

This comment has been revised 95 times.

[facebook-github-bot]

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

[vincentqb]

Can you also add the documentation? Some documentation tests are failing.

[facebook-github-bot]

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

[mthrok]

@vincentqb

I am confused about jit-ability of istft function. It seems that moving the function from torchaudio to pytorch slightly changes the property.
In both case, the bare istft function is jit-able, but if istft function is a part of custom function, it fails only when it's in pytorch.
The later case, the error message read, it requires implementation in aten.

import torch
import torchaudio.functional

# Toggle this to change the behavior.
istft = torch.istft
# istft = torchaudio.functional.istft

stft_kwargs = {
    'n_fft': 15,
    'hop_length': 3,
    'win_length': 11,
    'window': torch.hamming_window(11),
    'center': True,
    'pad_mode': 'constant',
    'normalized': True,
    'onesided': False,
}

istft_kwargs = stft_kwargs.copy()
istft_kwargs['length'] = None

input = torch.stft(torch.randn(3, 15), **stft_kwargs)


def nested_istft(
        stft_matrix, n_fft, hop_length, win_length, window, center, pad_mode, normalized, onesided, length):
    # type: (Tensor, int, Optional[int], Optional[int], Optional[Tensor], bool, str, bool, bool, Optional[int]) -> Tensor
    return istft(stft_matrix, n_fft, hop_length, win_length, window, center, pad_mode, normalized, onesided, length)


def test_func(istft_func):
    jit_istft = torch.jit.script(istft_func)
    py_result = istft_func(input, **istft_kwargs)
    jit_result = jit_istft(input, **istft_kwargs)
    assert torch.allclose(py_result, jit_result, atol=1e-6)


print('checking bare istft')
test_func(istft)
print('checking nested istft')
test_func(nested_istft)
print('ok')

When istft function is inside of torchaudio.functional, it works.

$ python test.py
checking bare istft
checking nested istft
ok

but when istft is inside of pytorch, then it fails.

$ python test.py
checking bare istft
checking nested istft
Traceback (most recent call last):
  File "foo.py", line 39, in <module>
    test_func(nested_istft)
  File "foo.py", line 30, in test_func
    jit_istft = torch.jit.script(istft_func)
  File "/scratch/moto/pytorch/torch/jit/__init__.py", line 1296, in script
    fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj))
RuntimeError:
Unknown builtin op: aten::istft.
Here are some suggestions:
	aten::stft
	aten::ifft
	aten::irfft

The original call is:
  File "foo.py", line 27
        stft_matrix, n_fft, hop_length, win_length, window, center, pad_mode, normalized, onesided, length):
    # type: (Tensor, int, Optional[int], Optional[int], Optional[Tensor], bool, str, bool, bool, Optional[int]) -> Tensor
    return istft(stft_matrix, n_fft, hop_length, win_length, window, center, pad_mode, normalized, onesided, length)
           ~~~~~ <--- HERE

[vincentqb]

This is related comment about aten error, and linked issue #30780.

[eellison]

@mthrok could you follow the example in #33737 ? Since the inputs all have one single type it shouldn't be very hard to make this torchscriptable. Just add
"istft" to ops = ["stft", "lu", "lu_unpack", "cdist"].

basically we generally resolve torch.{op} to "aten::{op}", but functions in torch/functional.py are an exception to this. There still exist some functions in torch/functional.py which we want to resolve to their aten equivalent like broadcast_tensors so ops have to opt-in to compiling the python function. (broadcast_tensors we didn't add bc we don't yet support (*args), but the aten builtin compiles correctly).

[mthrok]

Thanks @eellison for the suggestion. The trick did solve the issue above.

I applied that to PR, but now some tests are failing due to something related to ATen, which were not happening earlier.
On separate note I am working on porting istft to ATen, and I think ATen implementation will resolve this. So I will come back when I am done with ATen implementation of istft.

[facebook-github-bot]

@mthrok merged this pull request in 5a27ec0.