Audio augmentations with JAX

[niemiaszek]

I'm looking for some library for audio augmentations that could use JAX and nicely bind with tf.data pipelines.

I've seen AUX by DeepMind, which looks quite great, but has only one commit and limited capabilities compared to popular libraries (audiomentations/torchaudio).

I'm aware that audio is much less popular domain than LM/CV these days, but even looking at high amount of audio-related products developed by Google I'd expect some support for audio domain in JAX ecosystem.

I'm also wondering if I could find more information on audio pipelines at scale - f.e. what tools do people use to leverage HW and data? Would you still recommend offline dataset generation and storing pre-processed mel-spectrograms in TFDS?