torchaudio 2.7.1__cp312-cp312-win_amd64.whl → 2.9.0__cp312-cp312-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of torchaudio might be problematic. Click here for more details.

Files changed (92) hide show
  1. torchaudio/__init__.py +184 -33
  2. torchaudio/_extension/__init__.py +1 -14
  3. torchaudio/_extension/utils.py +0 -47
  4. torchaudio/_internal/module_utils.py +68 -10
  5. torchaudio/_torchcodec.py +340 -0
  6. torchaudio/datasets/cmuarctic.py +1 -1
  7. torchaudio/datasets/utils.py +1 -1
  8. torchaudio/functional/__init__.py +6 -3
  9. torchaudio/functional/_alignment.py +1 -1
  10. torchaudio/functional/filtering.py +70 -55
  11. torchaudio/functional/functional.py +31 -61
  12. torchaudio/lib/_torchaudio.pyd +0 -0
  13. torchaudio/lib/libtorchaudio.pyd +0 -0
  14. torchaudio/models/decoder/__init__.py +19 -1
  15. torchaudio/models/decoder/_ctc_decoder.py +6 -6
  16. torchaudio/models/decoder/_cuda_ctc_decoder.py +1 -1
  17. torchaudio/models/squim/objective.py +2 -2
  18. torchaudio/pipelines/_source_separation_pipeline.py +1 -1
  19. torchaudio/pipelines/_squim_pipeline.py +2 -2
  20. torchaudio/pipelines/_tts/utils.py +3 -1
  21. torchaudio/pipelines/rnnt_pipeline.py +4 -4
  22. torchaudio/transforms/__init__.py +4 -1
  23. torchaudio/transforms/_transforms.py +4 -3
  24. torchaudio/utils/__init__.py +2 -9
  25. torchaudio/utils/download.py +1 -1
  26. torchaudio/version.py +2 -2
  27. {torchaudio-2.7.1.dist-info → torchaudio-2.9.0.dist-info}/METADATA +15 -7
  28. torchaudio-2.9.0.dist-info/RECORD +85 -0
  29. {torchaudio-2.7.1.dist-info → torchaudio-2.9.0.dist-info}/top_level.txt +0 -1
  30. torchaudio/_backend/__init__.py +0 -61
  31. torchaudio/_backend/backend.py +0 -53
  32. torchaudio/_backend/common.py +0 -52
  33. torchaudio/_backend/ffmpeg.py +0 -334
  34. torchaudio/_backend/soundfile.py +0 -54
  35. torchaudio/_backend/soundfile_backend.py +0 -457
  36. torchaudio/_backend/sox.py +0 -91
  37. torchaudio/_backend/utils.py +0 -317
  38. torchaudio/backend/__init__.py +0 -8
  39. torchaudio/backend/_no_backend.py +0 -25
  40. torchaudio/backend/_sox_io_backend.py +0 -294
  41. torchaudio/backend/common.py +0 -13
  42. torchaudio/backend/no_backend.py +0 -14
  43. torchaudio/backend/soundfile_backend.py +0 -14
  44. torchaudio/backend/sox_io_backend.py +0 -14
  45. torchaudio/io/__init__.py +0 -13
  46. torchaudio/io/_effector.py +0 -347
  47. torchaudio/io/_playback.py +0 -72
  48. torchaudio/kaldi_io.py +0 -144
  49. torchaudio/prototype/__init__.py +0 -0
  50. torchaudio/prototype/datasets/__init__.py +0 -4
  51. torchaudio/prototype/datasets/musan.py +0 -67
  52. torchaudio/prototype/functional/__init__.py +0 -26
  53. torchaudio/prototype/functional/_dsp.py +0 -433
  54. torchaudio/prototype/functional/_rir.py +0 -379
  55. torchaudio/prototype/functional/functional.py +0 -190
  56. torchaudio/prototype/models/__init__.py +0 -36
  57. torchaudio/prototype/models/_conformer_wav2vec2.py +0 -794
  58. torchaudio/prototype/models/_emformer_hubert.py +0 -333
  59. torchaudio/prototype/models/conv_emformer.py +0 -525
  60. torchaudio/prototype/models/hifi_gan.py +0 -336
  61. torchaudio/prototype/models/rnnt.py +0 -711
  62. torchaudio/prototype/models/rnnt_decoder.py +0 -399
  63. torchaudio/prototype/pipelines/__init__.py +0 -12
  64. torchaudio/prototype/pipelines/_vggish/__init__.py +0 -3
  65. torchaudio/prototype/pipelines/_vggish/_vggish_impl.py +0 -233
  66. torchaudio/prototype/pipelines/_vggish/_vggish_pipeline.py +0 -82
  67. torchaudio/prototype/pipelines/hifigan_pipeline.py +0 -228
  68. torchaudio/prototype/pipelines/rnnt_pipeline.py +0 -58
  69. torchaudio/prototype/transforms/__init__.py +0 -9
  70. torchaudio/prototype/transforms/_transforms.py +0 -456
  71. torchaudio/sox_effects/__init__.py +0 -10
  72. torchaudio/sox_effects/sox_effects.py +0 -272
  73. torchaudio/utils/ffmpeg_utils.py +0 -11
  74. torchaudio/utils/sox_utils.py +0 -99
  75. torchaudio-2.7.1.dist-info/RECORD +0 -144
  76. torio/__init__.py +0 -8
  77. torio/_extension/__init__.py +0 -13
  78. torio/_extension/utils.py +0 -147
  79. torio/io/__init__.py +0 -9
  80. torio/io/_streaming_media_decoder.py +0 -978
  81. torio/io/_streaming_media_encoder.py +0 -502
  82. torio/lib/__init__.py +0 -0
  83. torio/lib/_torio_ffmpeg4.pyd +0 -0
  84. torio/lib/_torio_ffmpeg5.pyd +0 -0
  85. torio/lib/_torio_ffmpeg6.pyd +0 -0
  86. torio/lib/libtorio_ffmpeg4.pyd +0 -0
  87. torio/lib/libtorio_ffmpeg5.pyd +0 -0
  88. torio/lib/libtorio_ffmpeg6.pyd +0 -0
  89. torio/utils/__init__.py +0 -4
  90. torio/utils/ffmpeg_utils.py +0 -247
  91. {torchaudio-2.7.1.dist-info → torchaudio-2.9.0.dist-info}/WHEEL +0 -0
  92. {torchaudio-2.7.1.dist-info → torchaudio-2.9.0.dist-info}/licenses/LICENSE +0 -0
@@ -0,0 +1,340 @@
1
+ """TorchCodec integration for TorchAudio."""
2
+
3
+ import os
4
+ from typing import BinaryIO, Optional, Tuple, Union
5
+
6
+ import torch
7
+
8
+
9
+ def load_with_torchcodec(
10
+ uri: Union[BinaryIO, str, os.PathLike],
11
+ frame_offset: int = 0,
12
+ num_frames: int = -1,
13
+ normalize: bool = True,
14
+ channels_first: bool = True,
15
+ format: Optional[str] = None,
16
+ buffer_size: int = 4096,
17
+ backend: Optional[str] = None,
18
+ ) -> Tuple[torch.Tensor, int]:
19
+ """Load audio data from source using TorchCodec's AudioDecoder.
20
+
21
+ .. note::
22
+
23
+ This function supports the same API as :func:`~torchaudio.load`, and
24
+ relies on TorchCodec's decoding capabilities under the hood. It is
25
+ provided for convenience, but we do recommend that you port your code to
26
+ natively use ``torchcodec``'s ``AudioDecoder`` class for better
27
+ performance:
28
+ https://docs.pytorch.org/torchcodec/stable/generated/torchcodec.decoders.AudioDecoder.
29
+ As of TorchAudio 2.9, :func:`~torchaudio.load` relies on
30
+ :func:`~torchaudio.load_with_torchcodec`. Note that some parameters of
31
+ :func:`~torchaudio.load`, like ``normalize``, ``buffer_size``, and
32
+ ``backend``, are ignored by :func:`~torchaudio.load_with_torchcodec`.
33
+ To install torchcodec, follow the instructions at https://github.com/pytorch/torchcodec#installing-torchcodec.
34
+
35
+
36
+ Args:
37
+ uri (path-like object or file-like object):
38
+ Source of audio data. The following types are accepted:
39
+
40
+ * ``path-like``: File path or URL.
41
+ * ``file-like``: Object with ``read(size: int) -> bytes`` method.
42
+
43
+ frame_offset (int, optional):
44
+ Number of samples to skip before start reading data.
45
+ num_frames (int, optional):
46
+ Maximum number of samples to read. ``-1`` reads all the remaining samples,
47
+ starting from ``frame_offset``.
48
+ normalize (bool, optional):
49
+ TorchCodec always returns normalized float32 samples. This parameter
50
+ is ignored and a warning is issued if set to False.
51
+ Default: ``True``.
52
+ channels_first (bool, optional):
53
+ When True, the returned Tensor has dimension `[channel, time]`.
54
+ Otherwise, the returned Tensor's dimension is `[time, channel]`.
55
+ format (str or None, optional):
56
+ Format hint for the decoder. May not be supported by all TorchCodec
57
+ decoders. (Default: ``None``)
58
+ buffer_size (int, optional):
59
+ Not used by TorchCodec AudioDecoder. Provided for API compatibility.
60
+ backend (str or None, optional):
61
+ Not used by TorchCodec AudioDecoder. Provided for API compatibility.
62
+
63
+ Returns:
64
+ (torch.Tensor, int): Resulting Tensor and sample rate.
65
+ Always returns float32 tensors. If ``channels_first=True``, shape is
66
+ `[channel, time]`, otherwise `[time, channel]`.
67
+
68
+ Raises:
69
+ ImportError: If torchcodec is not available.
70
+ ValueError: If unsupported parameters are used.
71
+ RuntimeError: If TorchCodec fails to decode the audio.
72
+
73
+ Note:
74
+ - TorchCodec always returns normalized float32 samples, so the ``normalize``
75
+ parameter has no effect.
76
+ - The ``buffer_size`` and ``backend`` parameters are ignored.
77
+ - Not all audio formats supported by torchaudio backends may be supported
78
+ by TorchCodec.
79
+ """
80
+ # Import torchcodec here to provide clear error if not available
81
+ try:
82
+ from torchcodec.decoders import AudioDecoder
83
+ except ImportError as e:
84
+ raise ImportError(
85
+ "TorchCodec is required for load_with_torchcodec. " "Please install torchcodec to use this function."
86
+ ) from e
87
+
88
+ # Parameter validation and warnings
89
+ if not normalize:
90
+ import warnings
91
+
92
+ warnings.warn(
93
+ "TorchCodec AudioDecoder always returns normalized float32 samples. "
94
+ "The 'normalize=False' parameter is ignored.",
95
+ UserWarning,
96
+ stacklevel=2,
97
+ )
98
+
99
+ if buffer_size != 4096:
100
+ import warnings
101
+
102
+ warnings.warn("The 'buffer_size' parameter is not used by TorchCodec AudioDecoder.", UserWarning, stacklevel=2)
103
+
104
+ if backend is not None:
105
+ import warnings
106
+
107
+ warnings.warn("The 'backend' parameter is not used by TorchCodec AudioDecoder.", UserWarning, stacklevel=2)
108
+
109
+ if format is not None:
110
+ import warnings
111
+
112
+ warnings.warn("The 'format' parameter is not supported by TorchCodec AudioDecoder.", UserWarning, stacklevel=2)
113
+
114
+ # Create AudioDecoder
115
+ try:
116
+ decoder = AudioDecoder(uri)
117
+ except Exception as e:
118
+ raise RuntimeError(f"Failed to create AudioDecoder for {uri}: {e}") from e
119
+
120
+ # Get sample rate from metadata
121
+ sample_rate = decoder.metadata.sample_rate
122
+ if sample_rate is None:
123
+ raise RuntimeError("Unable to determine sample rate from audio metadata")
124
+
125
+ # Decode the entire file first, then subsample manually
126
+ # This is the simplest approach since torchcodec uses time-based indexing
127
+ try:
128
+ audio_samples = decoder.get_all_samples()
129
+ except Exception as e:
130
+ raise RuntimeError(f"Failed to decode audio samples: {e}") from e
131
+
132
+ data = audio_samples.data
133
+
134
+ # Apply frame_offset and num_frames (which are actually sample offsets)
135
+ if frame_offset > 0:
136
+ if frame_offset >= data.shape[1]:
137
+ # Return empty tensor if offset is beyond available data
138
+ empty_shape = (data.shape[0], 0) if channels_first else (0, data.shape[0])
139
+ return torch.zeros(empty_shape, dtype=torch.float32), sample_rate
140
+ data = data[:, frame_offset:]
141
+
142
+ if num_frames == 0:
143
+ # Return empty tensor if num_frames is 0
144
+ empty_shape = (data.shape[0], 0) if channels_first else (0, data.shape[0])
145
+ return torch.zeros(empty_shape, dtype=torch.float32), sample_rate
146
+ elif num_frames > 0:
147
+ data = data[:, :num_frames]
148
+
149
+ # TorchCodec returns data in [channel, time] format by default
150
+ # Handle channels_first parameter
151
+ if not channels_first:
152
+ data = data.transpose(0, 1) # [channel, time] -> [time, channel]
153
+
154
+ return data, sample_rate
155
+
156
+
157
+ def save_with_torchcodec(
158
+ uri: Union[str, os.PathLike],
159
+ src: torch.Tensor,
160
+ sample_rate: int,
161
+ channels_first: bool = True,
162
+ format: Optional[str] = None,
163
+ encoding: Optional[str] = None,
164
+ bits_per_sample: Optional[int] = None,
165
+ buffer_size: int = 4096,
166
+ backend: Optional[str] = None,
167
+ compression: Optional[Union[float, int]] = None,
168
+ ) -> None:
169
+ """Save audio data to file using TorchCodec's AudioEncoder.
170
+
171
+ .. note::
172
+
173
+ This function supports the same API as :func:`~torchaudio.save`, and
174
+ relies on TorchCodec's encoding capabilities under the hood. It is
175
+ provided for convenience, but we do recommend that you port your code to
176
+ natively use ``torchcodec``'s ``AudioEncoder`` class for better
177
+ performance:
178
+ https://docs.pytorch.org/torchcodec/stable/generated/torchcodec.encoders.AudioEncoder.
179
+ As of TorchAudio 2.9, :func:`~torchaudio.save` relies on
180
+ :func:`~torchaudio.save_with_torchcodec`. Note that some parameters of
181
+ :func:`~torchaudio.save`, like ``format``, ``encoding``,
182
+ ``bits_per_sample``, ``buffer_size``, and ``backend``, are ignored by
183
+ are ignored by :func:`~torchaudio.save_with_torchcodec`.
184
+ To install torchcodec, follow the instructions at https://github.com/pytorch/torchcodec#installing-torchcodec.
185
+
186
+ This function provides a TorchCodec-based alternative to torchaudio.save
187
+ with the same API. TorchCodec's AudioEncoder provides efficient encoding
188
+ with FFmpeg under the hood.
189
+
190
+ Args:
191
+ uri (path-like object):
192
+ Path to save the audio file. The file extension determines the format.
193
+
194
+ src (torch.Tensor):
195
+ Audio data to save. Must be a 1D or 2D tensor with float32 values
196
+ in the range [-1, 1]. If 2D, shape should be [channel, time] when
197
+ channels_first=True, or [time, channel] when channels_first=False.
198
+
199
+ sample_rate (int):
200
+ Sample rate of the audio data.
201
+
202
+ channels_first (bool, optional):
203
+ Indicates whether the input tensor has channels as the first dimension.
204
+ If True, expects [channel, time]. If False, expects [time, channel].
205
+ Default: True.
206
+
207
+ format (str or None, optional):
208
+ Audio format hint. Not used by TorchCodec (format is determined by
209
+ file extension). A warning is issued if provided.
210
+ Default: None.
211
+
212
+ encoding (str or None, optional):
213
+ Audio encoding. Not fully supported by TorchCodec AudioEncoder.
214
+ A warning is issued if provided. Default: None.
215
+
216
+ bits_per_sample (int or None, optional):
217
+ Bits per sample. Not directly supported by TorchCodec AudioEncoder.
218
+ A warning is issued if provided. Default: None.
219
+
220
+ buffer_size (int, optional):
221
+ Not used by TorchCodec AudioEncoder. Provided for API compatibility.
222
+ A warning is issued if not default value. Default: 4096.
223
+
224
+ backend (str or None, optional):
225
+ Not used by TorchCodec AudioEncoder. Provided for API compatibility.
226
+ A warning is issued if provided. Default: None.
227
+
228
+ compression (float, int or None, optional):
229
+ Compression level or bit rate. Maps to bit_rate parameter in
230
+ TorchCodec AudioEncoder. Default: None.
231
+
232
+ Raises:
233
+ ImportError: If torchcodec is not available.
234
+ ValueError: If input parameters are invalid.
235
+ RuntimeError: If TorchCodec fails to encode the audio.
236
+
237
+ Note:
238
+ - TorchCodec AudioEncoder expects float32 samples in [-1, 1] range.
239
+ - Some parameters (format, encoding, bits_per_sample, buffer_size, backend)
240
+ are not used by TorchCodec but are provided for API compatibility.
241
+ - The output format is determined by the file extension in the uri.
242
+ - TorchCodec uses FFmpeg under the hood for encoding.
243
+ """
244
+ # Import torchcodec here to provide clear error if not available
245
+ try:
246
+ from torchcodec.encoders import AudioEncoder
247
+ except ImportError as e:
248
+ raise ImportError(
249
+ "TorchCodec is required for save_with_torchcodec. " "Please install torchcodec to use this function."
250
+ ) from e
251
+
252
+ # Parameter validation and warnings
253
+ if format is not None:
254
+ import warnings
255
+
256
+ warnings.warn(
257
+ "The 'format' parameter is not used by TorchCodec AudioEncoder. "
258
+ "Format is determined by the file extension.",
259
+ UserWarning,
260
+ stacklevel=2,
261
+ )
262
+
263
+ if encoding is not None:
264
+ import warnings
265
+
266
+ warnings.warn(
267
+ "The 'encoding' parameter is not fully supported by TorchCodec AudioEncoder.", UserWarning, stacklevel=2
268
+ )
269
+
270
+ if bits_per_sample is not None:
271
+ import warnings
272
+
273
+ warnings.warn(
274
+ "The 'bits_per_sample' parameter is not directly supported by TorchCodec AudioEncoder.",
275
+ UserWarning,
276
+ stacklevel=2,
277
+ )
278
+
279
+ if buffer_size != 4096:
280
+ import warnings
281
+
282
+ warnings.warn("The 'buffer_size' parameter is not used by TorchCodec AudioEncoder.", UserWarning, stacklevel=2)
283
+
284
+ if backend is not None:
285
+ import warnings
286
+
287
+ warnings.warn("The 'backend' parameter is not used by TorchCodec AudioEncoder.", UserWarning, stacklevel=2)
288
+
289
+ # Input validation
290
+ if not isinstance(src, torch.Tensor):
291
+ raise ValueError(f"Expected src to be a torch.Tensor, got {type(src)}")
292
+
293
+ if src.dtype != torch.float32:
294
+ src = src.float()
295
+
296
+ if sample_rate <= 0:
297
+ raise ValueError(f"sample_rate must be positive, got {sample_rate}")
298
+
299
+ # Handle tensor shape and channels_first
300
+ if src.ndim == 1:
301
+ # Convert to 2D: [1, time] for channels_first=True
302
+ if channels_first:
303
+ data = src.unsqueeze(0) # [1, time]
304
+ else:
305
+ # For channels_first=False, input is [time] -> reshape to [time, 1] -> transpose to [1, time]
306
+ data = src.unsqueeze(1).transpose(0, 1) # [time, 1] -> [1, time]
307
+ elif src.ndim == 2:
308
+ if channels_first:
309
+ data = src # Already [channel, time]
310
+ else:
311
+ data = src.transpose(0, 1) # [time, channel] -> [channel, time]
312
+ else:
313
+ raise ValueError(f"Expected 1D or 2D tensor, got {src.ndim}D tensor")
314
+
315
+ # Create AudioEncoder
316
+ try:
317
+ encoder = AudioEncoder(data, sample_rate=sample_rate)
318
+ except Exception as e:
319
+ raise RuntimeError(f"Failed to create AudioEncoder: {e}") from e
320
+
321
+ # Determine bit_rate from compression parameter
322
+ bit_rate = None
323
+ if compression is not None:
324
+ if isinstance(compression, (int, float)):
325
+ bit_rate = int(compression)
326
+ else:
327
+ import warnings
328
+
329
+ warnings.warn(
330
+ f"Unsupported compression type {type(compression)}. "
331
+ "TorchCodec AudioEncoder expects int or float for bit_rate.",
332
+ UserWarning,
333
+ stacklevel=2,
334
+ )
335
+
336
+ # Save to file
337
+ try:
338
+ encoder.to_file(uri, bit_rate=bit_rate)
339
+ except Exception as e:
340
+ raise RuntimeError(f"Failed to save audio to {uri}: {e}") from e
@@ -129,7 +129,7 @@ class CMUARCTIC(Dataset):
129
129
  self._text = os.path.join(self._path, self._folder_text, self._file_text)
130
130
 
131
131
  with open(self._text, "r") as text:
132
- walker = csv.reader(text, delimiter="\n")
132
+ walker = csv.reader(text)
133
133
  self._walker = list(walker)
134
134
 
135
135
  def __getitem__(self, n: int) -> Tuple[Tensor, int, str, str]:
@@ -2,7 +2,7 @@ import logging
2
2
  import os
3
3
  import tarfile
4
4
  import zipfile
5
- from typing import Any, List, Optional
5
+ from typing import Any, List, Optional # noqa: F401
6
6
 
7
7
  import torchaudio
8
8
 
@@ -1,4 +1,6 @@
1
- from ._alignment import forced_align, merge_tokens, TokenSpan
1
+ from torchaudio._internal.module_utils import dropping_support
2
+
3
+ from ._alignment import forced_align as _forced_align, merge_tokens, TokenSpan
2
4
  from .filtering import (
3
5
  allpass_biquad,
4
6
  band_biquad,
@@ -23,11 +25,13 @@ from .filtering import (
23
25
  treble_biquad,
24
26
  vad,
25
27
  )
28
+
29
+ forced_align = dropping_support(_forced_align)
30
+
26
31
  from .functional import (
27
32
  add_noise,
28
33
  amplitude_to_DB,
29
34
  apply_beamforming,
30
- apply_codec,
31
35
  compute_deltas,
32
36
  convolve,
33
37
  create_dct,
@@ -106,7 +110,6 @@ __all__ = [
106
110
  "riaa_biquad",
107
111
  "treble_biquad",
108
112
  "vad",
109
- "apply_codec",
110
113
  "resample",
111
114
  "edit_distance",
112
115
  "pitch_shift",
@@ -70,7 +70,7 @@ def forced_align(
70
70
  assert target_lengths is not None
71
71
 
72
72
  paths, scores = torch.ops.torchaudio.forced_align(log_probs, targets, input_lengths, target_lengths, blank)
73
- return paths, scores
73
+ return paths, scores[:, torch.arange(scores.shape[1]), paths[0]]
74
74
 
75
75
 
76
76
  @dataclass
@@ -3,6 +3,7 @@ import warnings
3
3
  from typing import Optional
4
4
 
5
5
  import torch
6
+ import torch.nn.functional as F
6
7
  from torch import Tensor
7
8
 
8
9
  from torchaudio._extension import _IS_TORCHAUDIO_EXT_AVAILABLE
@@ -932,69 +933,83 @@ def _lfilter_core_generic_loop(input_signal_windows: Tensor, a_coeffs_flipped: T
932
933
 
933
934
 
934
935
  if _IS_TORCHAUDIO_EXT_AVAILABLE:
935
- _lfilter_core_cpu_loop = torch.ops.torchaudio._lfilter_core_loop
936
+ _lfilter_core_loop = torch.ops.torchaudio._lfilter_core_loop
936
937
  else:
937
- _lfilter_core_cpu_loop = _lfilter_core_generic_loop
938
+ _lfilter_core_loop = _lfilter_core_generic_loop
939
+
940
+
941
+ class DifferentiableFIR(torch.autograd.Function):
942
+ @staticmethod
943
+ def forward(ctx, waveform, b_coeffs):
944
+ n_order = b_coeffs.size(1)
945
+ n_channel = b_coeffs.size(0)
946
+ b_coeff_flipped = b_coeffs.flip(1).contiguous()
947
+ padded_waveform = F.pad(waveform, (n_order - 1, 0))
948
+ output = F.conv1d(padded_waveform, b_coeff_flipped.unsqueeze(1), groups=n_channel)
949
+ ctx.save_for_backward(waveform, b_coeffs, output)
950
+ return output
951
+
952
+ @staticmethod
953
+ def backward(ctx, dy):
954
+ x, b_coeffs, y = ctx.saved_tensors
955
+ n_batch = x.size(0)
956
+ n_channel = x.size(1)
957
+ n_order = b_coeffs.size(1)
958
+ db = (
959
+ F.conv1d(
960
+ F.pad(x, (n_order - 1, 0)).view(1, n_batch * n_channel, -1),
961
+ dy.view(n_batch * n_channel, 1, -1),
962
+ groups=n_batch * n_channel,
963
+ )
964
+ .view(n_batch, n_channel, -1)
965
+ .sum(0)
966
+ .flip(1)
967
+ if b_coeffs.requires_grad
968
+ else None
969
+ )
970
+ dx = F.conv1d(F.pad(dy, (0, n_order - 1)), b_coeffs.unsqueeze(1), groups=n_channel) if x.requires_grad else None
971
+ return (dx, db)
938
972
 
939
973
 
940
- def _lfilter_core(
941
- waveform: Tensor,
942
- a_coeffs: Tensor,
943
- b_coeffs: Tensor,
944
- ) -> Tensor:
974
+ class DifferentiableIIR(torch.autograd.Function):
975
+ @staticmethod
976
+ def forward(ctx, waveform, a_coeffs_normalized):
977
+ n_batch, n_channel, n_sample = waveform.shape
978
+ n_order = a_coeffs_normalized.size(1)
979
+ n_sample_padded = n_sample + n_order - 1
945
980
 
946
- if a_coeffs.size() != b_coeffs.size():
947
- raise ValueError(
948
- "Expected coeffs to be the same size."
949
- f"Found a_coeffs size: {a_coeffs.size()}, b_coeffs size: {b_coeffs.size()}"
981
+ a_coeff_flipped = a_coeffs_normalized.flip(1).contiguous()
982
+ padded_output_waveform = torch.zeros(
983
+ n_batch, n_channel, n_sample_padded, device=waveform.device, dtype=waveform.dtype
950
984
  )
951
- if waveform.ndim != 3:
952
- raise ValueError(f"Expected waveform to be 3 dimensional. Found: {waveform.ndim}")
953
- if not (waveform.device == a_coeffs.device == b_coeffs.device):
954
- raise ValueError(
955
- "Expected waveform and coeffs to be on the same device."
956
- f"Found: waveform device:{waveform.device}, a_coeffs device: {a_coeffs.device}, "
957
- f"b_coeffs device: {b_coeffs.device}"
985
+ _lfilter_core_loop(waveform, a_coeff_flipped, padded_output_waveform)
986
+ output = padded_output_waveform[:, :, n_order - 1 :]
987
+ ctx.save_for_backward(waveform, a_coeffs_normalized, output)
988
+ return output
989
+
990
+ @staticmethod
991
+ def backward(ctx, dy):
992
+ x, a_coeffs_normalized, y = ctx.saved_tensors
993
+ n_channel = x.size(1)
994
+ n_order = a_coeffs_normalized.size(1)
995
+ tmp = DifferentiableIIR.apply(dy.flip(2).contiguous(), a_coeffs_normalized).flip(2)
996
+ dx = tmp if x.requires_grad else None
997
+ da = (
998
+ -(
999
+ tmp.transpose(0, 1).reshape(n_channel, 1, -1)
1000
+ @ F.pad(y, (n_order - 1, 0)).unfold(2, n_order, 1).transpose(0, 1).reshape(n_channel, -1, n_order)
1001
+ )
1002
+ .squeeze(1)
1003
+ .flip(1)
1004
+ if a_coeffs_normalized.requires_grad
1005
+ else None
958
1006
  )
1007
+ return (dx, da)
959
1008
 
960
- n_batch, n_channel, n_sample = waveform.size()
961
- n_order = a_coeffs.size(1)
962
- if n_order <= 0:
963
- raise ValueError(f"Expected n_order to be positive. Found: {n_order}")
964
-
965
- # Pad the input and create output
966
1009
 
967
- padded_waveform = torch.nn.functional.pad(waveform, [n_order - 1, 0])
968
- padded_output_waveform = torch.zeros_like(padded_waveform)
969
-
970
- # Set up the coefficients matrix
971
- # Flip coefficients' order
972
- a_coeffs_flipped = a_coeffs.flip(1)
973
- b_coeffs_flipped = b_coeffs.flip(1)
974
-
975
- # calculate windowed_input_signal in parallel using convolution
976
- input_signal_windows = torch.nn.functional.conv1d(padded_waveform, b_coeffs_flipped.unsqueeze(1), groups=n_channel)
977
-
978
- input_signal_windows.div_(a_coeffs[:, :1])
979
- a_coeffs_flipped.div_(a_coeffs[:, :1])
980
-
981
- if (
982
- input_signal_windows.device == torch.device("cpu")
983
- and a_coeffs_flipped.device == torch.device("cpu")
984
- and padded_output_waveform.device == torch.device("cpu")
985
- ):
986
- _lfilter_core_cpu_loop(input_signal_windows, a_coeffs_flipped, padded_output_waveform)
987
- else:
988
- _lfilter_core_generic_loop(input_signal_windows, a_coeffs_flipped, padded_output_waveform)
989
-
990
- output = padded_output_waveform[:, :, n_order - 1 :]
991
- return output
992
-
993
-
994
- if _IS_TORCHAUDIO_EXT_AVAILABLE:
995
- _lfilter = torch.ops.torchaudio._lfilter
996
- else:
997
- _lfilter = _lfilter_core
1010
+ def _lfilter(waveform, a_coeffs, b_coeffs):
1011
+ filtered_waveform = DifferentiableFIR.apply(waveform, b_coeffs / a_coeffs[:, 0:1])
1012
+ return DifferentiableIIR.apply(filtered_waveform, a_coeffs / a_coeffs[:, 0:1])
998
1013
 
999
1014
 
1000
1015
  def lfilter(waveform: Tensor, a_coeffs: Tensor, b_coeffs: Tensor, clamp: bool = True, batching: bool = True) -> Tensor: