Add autograd test to T.TimeStretch (and F.phase_vocoder)

[mthrok]

UPDATE
Due to the conversion from cartesian coordinate to polar coordinate, F.phase_vocoder (and T.TimeStretch) is not differentiable everywhere.

---

Part of #1337
Looks like rate=0.7 is not differentiable to second degree.
Let's see what are the value ranges of rate, that are differentiable to second degree.
T.TimeStretch manipulates spectrogram in polar coordinate.
The angle conversion from Cartesian to Polar (atan2(Img, Real)) does not have a
nice differentiability around (x=0, y=0) and theta = +/- pi, therefore T.TimeStrech is not differentiable there.

TODO

• Check if rate=0.7 issue happens without complex Tensor adaptation ([TEST] Autograd phase vocoder ref #1430)
  -> Yes it does. In fact the adaptation of complex Tensor makes the TimeStretch work with rate in range 0.8 ~ 1.3.
  -> Filed the issue for rate=0.7 Gradcheck failure for torchaudio.functional.phase_vocoder pytorch#55557

[mthrok]

@anjali411 I have removed the failing case 0.7. CR please.

[mthrok]

Hmm, now the tests are failing for all the cases....

[anjali411]

this issue looks similar to pytorch/pytorch#42952, only that the function to blame here is likely angle which calls into std::arg which uses atan2: https://github.com/llvm/llvm-project/blob/ce8f10d6cbe03c1a78a7938070654d363e25f27d/libcxx/include/complex#L860-L866. This also explains why the autograd tests were failing on the previous version using pseudo complex types since we still call into atan2 via the angle function defined in TorchAudio here:

audio/torchaudio/functional/functional.py

Lines 540 to 551 in 0fbfca5

	def angle(
	complex_tensor: Tensor
	) -> Tensor:
	r"""Compute the angle of complex tensor input.
	Args:
	complex_tensor (Tensor): Tensor shape of `(..., complex=2)`
	Return:
	Tensor: Angle of a complex tensor. Shape of `(..., )`
	"""
	return torch.atan2(complex_tensor[..., 1], complex_tensor[..., 0])

never mind I was wrong, as @lezcano mentions here, that atan2 is not differentiable at 0 and so we shouldn't gradcheck here. @mthrok One thing to look into here is if the phase_vocoder function is expected to be differentiable overall, in which case, we might wanna define custom gradients.

[mthrok]

@mthrok One thing to look into here is if the phase_vocoder function is expected to be differentiable overall, in which case, we might wanna define custom gradients.

@anjali411 That's a tough question. I am not sure if it is supposed to be a differential. The primal profit of autograd support for ops in torchaudio is that it might enable the E2E task of combined training. For that first-order gradient is enough.

[mthrok]

@anjali411

My brief understanding of F.phase_vocoder algorithm is that it performs phase shift of complex spectrogram in polar coordinate, so it is inevitable to use atan when converting the input complex value (in cartesian form) to polar form.

So, unless there is another formula, I say that F.phase_vocoder and T.TimeStretch do not support autograd, and it has nothing to do with native complex Tensor.

I am going to close this PR.

[lezcano]

For reference, there's no other formula for the gradient. See this plot for what's going on at zero:
https://www.wolframalpha.com/input/?i=plot+atan2%28x%2Cy%29

It would be enough to make sure that the tests do not test points near zero, but I do not know whether that's possible in this particular use case

[mthrok]

Oh, what I meant by "another formula" is a formula for achieving the same result as the current F.phase_vocoder without converting the input spectrogram to polar coordinate.

Since the angle is applied to the input Tensor, I do not think it's practically possible to ensure non-zero input for the real use case, even if we can do so for testing.

[mthrok]

It would be enough to make sure that the tests do not test points near zero, but I do not know whether that's possible in this particular use case

@lezcano

Thanks for the advice. Following @anjali411's suggestion I tried to see if it's possible to come up with a guideline for workaround the zero values. No luck, I bet someone has already invented it if it's something I can come up with easily.

@Anjali

I gave some thoughts on if it is possible to work around the in-differentiability. I think, if we really want to tackle this, we need to re-invent arctangent function.

I think the partial derivatives can be expressed as
plot
plot

And they go to as close as +/- infinity around zero depending on the phase.
For the sake of T.TimeStretch, it's okay to have zero gradient around the origin because in this context, abs(z)==0 means that the input signal has only an negligible amount of energy in the corresponding frequency bin and for the sake of NN model training passing 0 gradient makes sense. Now the problem is how we can invent a formula that has a nice differentiability around 0 but gets closer to atan2 as the input gets far from the origin.

Essentially this function needs to satisfy some properties like

and if z is not close to the origin,

and they connect soothly

similarly their derivatives have to meet the properties like

but have the regular partial derivatives of atan2 when z is far from 0, while they connect nicely.

[lezcano]

I really like the name super_atan2, not gonna lie :D

Now, the first thing that comes to mind is to interpolate with a function that decreases very quickly to zero. The typical example of such a function is e^{-1/x}. So, if we denote r = (x²+y²)/α², for a small α > 0, we could consider
super_atan2(y,x) = (1-r)e^{-1/r} + r atan2(y,x) on the ball of radius α x²+y² < α², and then we would like to extend it to be atan2(y,x) outside of this ball. This extension is continuous but not differentiable, but that should be fine, I reckon.

Now, if we were to do this, it might be even simpler to just have a cut-off at some small α for the derivative, and declare it to be zero after that.

Now, all these things are clearly just hacks. The real question would then be: is it really necessary to make this function differentiable, or would it be better to leave it as it is, and then react depending on whether people file issues?

[mthrok]

@lezcano

Now, the first thing that comes to mind is to interpolate with a function that decreases very quickly to zero. The typical example of such a function is e^{-1/x}. So, if we denote r = (x²+y²)/α², for a small α > 0, we could consider
super_atan2(y,x) = (1-r)e^{-1/r} + r atan2(y,x) on the ball of radius α x²+y² < α², and then we would like to extend it to be atan2(y,x) outside of this ball. This extension is continuous but not differentiable, but that should be fine, I reckon.

Now, if we were to do this, it might be even simpler to just have a cut-off at some small α for the derivative, and declare it to be zero after that.

wow those are cool options! Thanks for the suggestions.

Now, all these things are clearly just hacks. The real question would then be: is it really necessary to make this function differentiable, or would it be better to leave it as it is, and then react depending on whether people file issues?

Well, for the scope of native complex Tensor support, we do not need this. We can just say TimeStretch does not support autograd.
But this got me thinking, it's not just this op but any op that manipulates data in polar coordinate is not differentiable. In case of signal processing and spectrogram, having 0 is inevitable, so the regular atan makes it impossible to use polar coordinate in training pipeline, and having a workaround might open a new door for research. Of course I have no idea it will be useful or not.

At this moment it's more of a thought experiment, but this could be also an interesting research topic.

[lezcano]

What I would suggest is, rather than saying that it's not differentiable, I would write down the differentiability issue in a .. warning::. That's what we are doing in torch.linalg, where it's much less obvious when do the gradients explode. See for example the penultimate warning in https://pytorch.org/docs/master/linalg.html?highlight=linalg%20svd#torch.linalg.svd

[mthrok]

Thanks for the tips. Will do so.

[mthrok]

@anjali411 Following the tips from @lezcano, I realized that we need a test to claim the differentiability at non-zero point. So I re-opened the PR. (sorry for flip-flopping)

In the new approach, I push the points close to zero away from zero. so that all the complex value elements have norm at least epsilon.