Repository: pytorch/audio

Issue:
Support multiple filters in lfilter
## 🚀 Feature
It would be nice if we can apply multiple filter in `F.lfilter`

## Motivation

https://github.com/pytorch/audio/issues/1517#issuecomment-845125829

## Pitch

Allow [the `a_coeffs` and `b_coefds` parameters of `F.lfilter`](https://github.com/pytorch/audio/blob/14dd917ec60fa69ce3f7c6e3f2eaf520e67928b5/torchaudio/functional/filtering.py#L879-L883) to be 2D Tensors.

## Alternatives

Currently one has to use for-loop for achieving this.

## Additional context

There is also a feature proposal for extending `lfilter` for batch processing. https://github.com/pytorch/audio/issues/1476


Discussion:

Resolving patch:
diff --git a/torchaudio/functional/filtering.py b/torchaudio/functional/filtering.py
--- a/torchaudio/functional/filtering.py
+++ b/torchaudio/functional/filtering.py
@@ -857,13 +857,14 @@ def highpass_biquad(
 
 
 def _lfilter_core_generic_loop(input_signal_windows: Tensor, a_coeffs_flipped: Tensor, padded_output_waveform: Tensor):
-    n_order = a_coeffs_flipped.size(0)
-    for i_sample, o0 in enumerate(input_signal_windows.t()):
+    n_order = a_coeffs_flipped.size(1)
+    a_coeffs_flipped = a_coeffs_flipped.unsqueeze(2)
+    for i_sample, o0 in enumerate(input_signal_windows.permute(2, 0, 1)):
         windowed_output_signal = padded_output_waveform[
-            :, i_sample:i_sample + n_order
+            :, :, i_sample:i_sample + n_order
         ]
-        o0.addmv_(windowed_output_signal, a_coeffs_flipped, alpha=-1)
-        padded_output_waveform[:, i_sample + n_order - 1] = o0
+        o0 -= (windowed_output_signal.transpose(0, 1) @ a_coeffs_flipped)[..., 0].t()
+        padded_output_waveform[:, :, i_sample + n_order - 1] = o0
 
 
 try:
@@ -879,13 +880,13 @@ def _lfilter_core(
     b_coeffs: Tensor,
 ) -> Tensor:
 
-    assert a_coeffs.size(0) == b_coeffs.size(0)
-    assert len(waveform.size()) == 2
+    assert a_coeffs.size() == b_coeffs.size()
+    assert len(waveform.size()) == 3
     assert waveform.device == a_coeffs.device
     assert b_coeffs.device == a_coeffs.device
 
-    n_channel, n_sample = waveform.size()
-    n_order = a_coeffs.size(0)
+    n_batch, n_channel, n_sample = waveform.size()
+    n_order = a_coeffs.size(1)
     assert n_order > 0
 
     # Pad the input and create output
@@ -895,17 +896,18 @@ def _lfilter_core(
 
     # Set up the coefficients matrix
     # Flip coefficients' order
-    a_coeffs_flipped = a_coeffs.flip(0)
-    b_coeffs_flipped = b_coeffs.flip(0)
+    a_coeffs_flipped = a_coeffs.flip(1)
+    b_coeffs_flipped = b_coeffs.flip(1)
 
     # calculate windowed_input_signal in parallel using convolution
     input_signal_windows = torch.nn.functional.conv1d(
-        padded_waveform.unsqueeze(1),
-        b_coeffs_flipped.view(1, 1, -1)
-    ).squeeze(1)
+        padded_waveform,
+        b_coeffs_flipped.unsqueeze(1),
+        groups=n_channel
+    )
 
-    input_signal_windows.div_(a_coeffs[0])
-    a_coeffs_flipped.div_(a_coeffs[0])
+    input_signal_windows.div_(a_coeffs[:, :1])
+    a_coeffs_flipped.div_(a_coeffs[:, :1])
 
     if input_signal_windows.device == torch.device('cpu') and\
        a_coeffs_flipped.device == torch.device('cpu') and\
@@ -914,9 +916,10 @@ def _lfilter_core(
     else:
         _lfilter_core_generic_loop(input_signal_windows, a_coeffs_flipped, padded_output_waveform)
 
-    output = padded_output_waveform[:, n_order - 1:]
+    output = padded_output_waveform[:, :, n_order - 1:]
     return output
 
+
 try:
     _lfilter = torch.ops.torchaudio._lfilter
 except RuntimeError as err:
@@ -938,21 +941,32 @@ def lfilter(
 
     Args:
         waveform (Tensor): audio waveform of dimension of ``(..., time)``.  Must be normalized to -1 to 1.
-        a_coeffs (Tensor): denominator coefficients of difference equation of dimension of ``(n_order + 1)``.
+        a_coeffs (Tensor): denominator coefficients of difference equation of dimension of either
+                                1D with shape ``(num_order + 1)`` or 2D with shape ``(num_filters, num_order + 1)``.
                                 Lower delays coefficients are first, e.g. ``[a0, a1, a2, ...]``.
                                 Must be same size as b_coeffs (pad with 0's as necessary).
-        b_coeffs (Tensor): numerator coefficients of difference equation of dimension of ``(n_order + 1)``.
-                                 Lower delays coefficients are first, e.g. ``[b0, b1, b2, ...]``.
-                                 Must be same size as a_coeffs (pad with 0's as necessary).
+        b_coeffs (Tensor): numerator coefficients of difference equation of dimension of either
+                                1D with shape ``(num_order + 1)`` or 2D with shape ``(num_filters, num_order + 1)``.
+                                Lower delays coefficients are first, e.g. ``[b0, b1, b2, ...]``.
+                                Must be same size as a_coeffs (pad with 0's as necessary).
         clamp (bool, optional): If ``True``, clamp the output signal to be in the range [-1, 1] (Default: ``True``)
 
     Returns:
-        Tensor: Waveform with dimension of ``(..., time)``.
+        Tensor: Waveform with dimension of either ``(..., num_filters, time)`` if ``a_coeffs`` and ``b_coeffs``
+                are 2D Tensors, or ``(..., time)`` otherwise.
     """
+    assert a_coeffs.size() == b_coeffs.size()
+    assert a_coeffs.ndim <= 2
+
+    if a_coeffs.ndim > 1:
+        waveform = torch.stack([waveform] * a_coeffs.shape[0], -2)
+    else:
+        a_coeffs = a_coeffs.unsqueeze(0)
+        b_coeffs = b_coeffs.unsqueeze(0)
+
     # pack batch
     shape = waveform.size()
-    waveform = waveform.reshape(-1, shape[-1])
-
+    waveform = waveform.reshape(-1, a_coeffs.shape[0], shape[-1])
     output = _lfilter(waveform, a_coeffs, b_coeffs)
 
     if clamp:


Test patch:
diff --git a/test/torchaudio_unittest/functional/autograd_impl.py b/test/torchaudio_unittest/functional/autograd_impl.py
--- a/test/torchaudio_unittest/functional/autograd_impl.py
+++ b/test/torchaudio_unittest/functional/autograd_impl.py
@@ -59,6 +59,15 @@ def test_lfilter_all_inputs(self):
         b = torch.tensor([0.4, 0.2, 0.9])
         self.assert_grad(F.lfilter, (x, a, b))
 
+    def test_lfilter_filterbanks(self):
+        torch.random.manual_seed(2434)
+        x = get_whitenoise(sample_rate=22050, duration=0.01, n_channels=3)
+        a = torch.tensor([[0.7, 0.2, 0.6],
+                          [0.8, 0.2, 0.9]])
+        b = torch.tensor([[0.4, 0.2, 0.9],
+                          [0.7, 0.2, 0.6]])
+        self.assert_grad(F.lfilter, (x, a, b))
+
     def test_biquad(self):
         torch.random.manual_seed(2434)
         x = get_whitenoise(sample_rate=22050, duration=0.01, n_channels=1)
diff --git a/test/torchaudio_unittest/functional/functional_impl.py b/test/torchaudio_unittest/functional/functional_impl.py
--- a/test/torchaudio_unittest/functional/functional_impl.py
+++ b/test/torchaudio_unittest/functional/functional_impl.py
@@ -67,18 +67,22 @@ def test_lfilter_clamp(self):
         assert output_signal.max() > 1
 
     @parameterized.expand([
-        ((44100,),),
-        ((3, 44100),),
-        ((2, 3, 44100),),
-        ((1, 2, 3, 44100),)
+        ((44100,), (4,), (44100,)),
+        ((3, 44100), (4,), (3, 44100,)),
+        ((2, 3, 44100), (4,), (2, 3, 44100,)),
+        ((1, 2, 3, 44100), (4,), (1, 2, 3, 44100,)),
+        ((44100,), (2, 4), (2, 44100)),
+        ((3, 44100), (1, 4), (3, 1, 44100)),
+        ((1, 2, 44100), (3, 4), (1, 2, 3, 44100))
     ])
-    def test_lfilter_shape(self, shape):
+    def test_lfilter_shape(self, input_shape, coeff_shape, target_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)
+        waveform = torch.rand(*input_shape, dtype=self.dtype, device=self.device)
+        b_coeffs = torch.rand(*coeff_shape, dtype=self.dtype, device=self.device)
+        a_coeffs = torch.rand(*coeff_shape, dtype=self.dtype, device=self.device)
         output_waveform = F.lfilter(waveform, a_coeffs, b_coeffs)
-        assert shape == waveform.size() == output_waveform.size()
+        assert input_shape == waveform.size()
+        assert target_shape == output_waveform.size()
 
     def test_lfilter_9th_order_filter_stability(self):
         """