Faustax

Faustax supplies differentiable audio effects for JAX. Each effect is a Faust program. The Faust compiler converts each program ahead of time into a Flax NNX module.

Faustax is the JAX equivalent of dasp-pytorch. It supplies batched, differentiable audio processors for parameter estimation, style transfer, automatic mixing, and data augmentation. Faustax covers every processor that dasp-pytorch implements. The implementation is different. dasp-pytorch approximates IIR filters in the frequency domain. In Faustax, each effect is a small Faust program. Most of these programs use the mature and widely used stdfaust libraries. The NNX backend of Faust compiles each program into an exact per-sample recurrence. This recurrence is compatible with jax.jit and jax.vmap.

Faustax and dasp-pytorch have different performance characteristics because the implementations are different. The exact per-sample recurrence of Faustax is sequential in time. The frequency-domain approximations of dasp-pytorch are a few large batched tensor operations. On the CPU, the Faustax EQ costs approximately 4x the dasp-pytorch forward time and 6-8x the dasp-pytorch gradient time. This cost is the cost of coefficient-exact IIR output. The Faustax compressor is at parity with the dasp-pytorch compressor. The Faustax FFT-convolution reverb trains 20-30x faster than the direct convolution of dasp-pytorch. On the GPU, a sequential scan costs approximately 0.5-1 s for each call at any batch size. Thus dasp-pytorch is faster at small batch sizes. The Faustax wall time stays almost constant as the batch grows. Run the Faustax scan on the CPU when you have fewer than approximately 64 parallel lanes. Use Faustax when you need exact output or parity between training and streaming deployment. Also use Faustax for fast reverb and dynamics training, or for wide-batch throughput on the CPU. Use frequency sampling, as in dasp-pytorch, when small-batch EQ gradients on the GPU are the most important cost. The Performance page contains the measured tables for each claim in this paragraph.

End users do not need the Faust compiler. The repository contains the generated Python modules. The only runtime dependencies are jax, flax, numpy, librosa, and safetensors.