standardizing freq/time axis

[vincentqb]

Let's use the same convention between masking functions (lines 826 and 862), while also aligning with the batching convention we use (lines 823 and 829), see #391.

Reference: #285

[vincentqb]

Test:

diff --git a/torchaudio/functional.py b/torchaudio/functional.py
index 2690d04..22e2ddc 100644
--- a/torchaudio/functional.py
+++ b/torchaudio/functional.py
@@ -829,8 +829,8 @@ def mask_along_axis_iid(specgrams, mask_param, mask_value, axis):
         torch.Tensor: Masked spectrograms of dimensions (..., freq, time)
     """
 
-    if axis != -2 and axis != -1:
-        raise ValueError('Only Frequency and Time masking are supported')
+    # if axis != -2 and axis != -1:
+    #    raise ValueError('Only Frequency and Time masking are supported')
 
     value = torch.rand(specgrams.shape[:2]) * mask_param
     min_value = torch.rand(specgrams.shape[:2]) * (specgrams.size(axis) - value)
@@ -876,9 +876,9 @@ def mask_along_axis(specgram, mask_param, mask_value, axis):
     mask_end = (min_value.long() + value.long()).squeeze()
 
     assert mask_end - mask_start < mask_param
-    if axis == -2:
+    if axis == -2 or axis == 1:
         specgram[:, mask_start:mask_end] = mask_value
-    elif axis == -1:
+    elif axis == -1 or axis == 2:
         specgram[:, :, mask_start:mask_end] = mask_value
     else:
         raise ValueError('Only Frequency and Time masking are supported')

import torch, torchaudio

torch.random.manual_seed(42)
specgram = torch.randn(1, 2, 5, 7)

torch.random.manual_seed(42)
output1 = torchaudio.functional.mask_along_axis_iid(specgram.clone(), mask_param=5, mask_value=0, axis=-1)

torch.random.manual_seed(42)
output2 = torchaudio.functional.mask_along_axis_iid(specgram.clone(), mask_param=5, mask_value=0, axis=3)

print(torch.allclose(output1, output2))

torch.random.manual_seed(42)
output1 = torchaudio.functional.mask_along_axis_iid(specgram.clone(), mask_param=5, mask_value=0, axis=-2)

torch.random.manual_seed(42)
output2 = torchaudio.functional.mask_along_axis_iid(specgram.clone(), mask_param=5, mask_value=0, axis=2)

print(torch.allclose(output1, output2))

torch.random.manual_seed(42)
output1 = torchaudio.functional.mask_along_axis(specgram.clone(), mask_param=5, mask_value=0, axis=2)

torch.random.manual_seed(42)
output2 = torchaudio.functional.mask_along_axis(specgram.clone(), mask_param=5, mask_value=0, axis=-1)

print(torch.allclose(output1, output2))

torch.random.manual_seed(42)
output1 = torchaudio.functional.mask_along_axis(specgram.clone(), mask_param=5, mask_value=0, axis=1)

torch.random.manual_seed(42)
output2 = torchaudio.functional.mask_along_axis(specgram.clone(), mask_param=5, mask_value=0, axis=-2)

print(torch.allclose(output1, output2))

True
True
True
True

[vincentqb]

@mthrok -- thoughts on this convention?

[mthrok]

Looks good.

[mthrok]

I feel it a little bit too strict to disallow positive indices, but as long as it is explained in doctoring, I think it's fair.

[mthrok]

oh wait, isn't this BC breaking?

[mthrok]

nit

Suggested change

	if axis != -2 and axis != -1:
	if axis not in [-1, -2]:

[vincentqb]

Unless there is a request for this, let's not do this BC-breaking change.