Add test for validating lfilter shape (#1360)

pytorch · Mar 5, 2021 · e868d24 · e868d24
1 parent 7ab0820
commit e868d24
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/test/torchaudio_unittest/functional/functional_impl.py b/test/torchaudio_unittest/functional/functional_impl.py
@@ -31,6 +31,20 @@ def test_clamp(self):
         output_signal = F.lfilter(input_signal, a_coeffs, b_coeffs, clamp=False)
         assert output_signal.max() > 1
 
+    @parameterized.expand([
+        ((44100,),),
+        ((3, 44100),),
+        ((2, 3, 44100),),
+        ((1, 2, 3, 44100),)
+    ])
+    def test_shape(self, shape):
+        torch.random.manual_seed(42)
+        waveform = torch.rand(*shape, dtype=self.dtype, device=self.device)
+        b_coeffs = torch.tensor([0, 0, 0, 1], dtype=self.dtype, device=self.device)
+        a_coeffs = torch.tensor([1, 0, 0, 0], dtype=self.dtype, device=self.device)
+        output_waveform = F.lfilter(waveform, a_coeffs, b_coeffs)
+        assert shape == waveform.size() == output_waveform.size()
+
 
 class Spectrogram(common_utils.TestBaseMixin):
     @parameterized.expand([(0., ), (1., ), (2., ), (3., )])