Migrate TorchAudio to use complex tensors in place of real tensors of shape (..., 2)

[anjali411]

Complex Tensors were released in 1.6 release of torch. The progress of complex numbers support in torch can be tracked here. This issue is meant to outline the tasks that need to be completed to migrate torchaudio API to use these newly added complex tensors.

The general plan for this migration, would be to add a local USE_COMPLEX flag to avoid bc breaking changes. This flag detect if the input is complex or not. If it's not complex, the old code (that supports the real tensors with shape (..., 2)) will be used and the returned tensor will be a real tensor. If it's complex, the returned tensor will be a complex tensor.
Here's a demo PR for this migration: #758

Here's the tracker for the issue for deprecation of torch.fft functions. The new torch.fft.fft functions will take as input and return complex tensors instead of the real representation of complex tensors.

torchaudio.functional

• spectrogram (uses torch.stft)

• (deprecate) complex_norm (same as torch.abs())

• (deprecate) mag_phase

• (deprecate) angle (same astorch.angle())

• [WIP] phase_vocoder

• _measure (uses torch.rfft, complex_norm) (Update torch.rfft to torch.fft.rfft and complex tensor #941)

• griffinlim (uses torch.istft, torch.stft)

torchaudio.transforms

• Spectrogram

• MelSpectrogram

• GriffinLim

• [WIP] TimeStretch

• (deprecate) ComplexNorm

• Vad (Update torch.rfft to torch.fft.rfft and complex tensor #941)

torchaudio.compliance.kaldi

• spectrogram (uses torch.rfft) (Update torch.rfft to torch.fft.rfft and complex tensor #941)

• bank (uses torch.rfft) (Update torch.rfft to torch.fft.rfft and complex tensor #941)

cc. @vincentqb @mthrok
cc. @mruberry, who is working on the old API of torch.ftt deprecation, for visibility.

[mthrok]

Thanks @anjali411.

Do we have an estimate on completion of this migration?
(By the completion of the migration, I mean we will only support complex dtype.)
I assume after all the implementation is completed, we want to have certain release cycles where users can opt-in for complex dtype, and make sure that some major downstream users migrate too.

[vincentqb]

The general plan for this migration, would be to add a local USE_COMPLEX flag to avoid bc breaking changes. This flag detect if the input is complex or not. If it's not complex, the old code (that supports the real tensors with shape (..., 2)) will be used and the returned tensor will be a real tensor. If it's complex, the returned tensor will be a complex tensor.

To elaborate, the plan we had discussed was to:

• We need to decide whether torchaudio needs to test that autograd runs without errors with complex numbers.
• For each transform having a complex equivalent, detect when a complex dtype tensor is offered (edited following comment below) add use_complex flag, and use the complex implementation in this case.
• Once all operations with complex support are implemented, we add a deprecation warning to all transform for at least one release that triggers when the non-complex path is used.
• After at least one release, we then remove the code path for non-complex tensor.

Can you remind me why the flag USE_COMPLEX is needed here @anjali411 ?

[anjali411]

Can you remind me why the flag USE_COMPLEX is needed here @anjali411 ?

to detect if the input is complex and if we should use the code for complex case (instead of the (..., 2) style float tensors. It also ensures we don't break autograd support in torchaudio (since it relies on torch for autograd correctness).

[vincentqb]

Can you remind me why the flag USE_COMPLEX is needed here @anjali411 ?
to detect if the input is complex and if we should use the code for complex case (instead of the (..., 2) style float tensors. It also ensures we don't break autograd support in torchaudio (since it relies on torch for autograd correctness).

Right, so the concern is that we could be passing a real tensor, which we want to treat as complex with zero imaginary for instance. In that case, we no longer need to detect whether the tensor is real or complex, right? We can just go through the complex path. Is that right?

[mthrok]

This issue is replaced by #1337