Custom gradient primitives

The generated modules differentiate their per-sample scan with standard reverse-mode autodiff (BPTT). BPTT stores the scan carry at each step. A better rule exists for recursive filters. The rule comes from Chin-Yun Yu’s torchlpc / torchcomp (DAFx 2024, arXiv:2404.07970). The vector-Jacobian product of an all-pole filter is the same filter, run backwards in time with time-shifted coefficients. The rule adds a per-sample outer product. Thus the backward pass needs only the output of the filter: one array, not the carry history. The backward pass also reuses the forward kernel. Thus higher-order derivatives compose automatically. faustax.ops ports these rules to JAX as jax.custom_vjp primitives. The primitives are unbatched, and jax.vmap adds batching:

• ops.allpole(x, a, zi) — the time-varying all-pole filter y[t] = x[t] − Σₖ a[t,k]·y[t−k−1]. It is the JAX equivalent of torchlpc’s sample_wise_lpc (see also jaxpole).

• ops.linear_recurrence(decay, x, zi) — a first-order convenience wrapper.

• ops.ballistics(x, at, rt, zi) — the switched attack/release one-pole smoother (torchcomp’s compressor_core). The backward pass keeps the per-sample switch decision constant. This behavior matches torchcomp’s piecewise treatment.

• ops.fir(b, x, zi) and ops.lfilter(b, a, x, zi) — time-varying (or constant) FIR and direct-form-II IIR filters. They follow philtorch’s exact DF2 decomposition: lfilter = fir(b, allpole(x, a)) over one shared delay line. Thus the recursive part uses the custom VJP. With constant coefficients and zero state, the result matches scipy.signal.lfilter (a excludes the leading 1). return_zf=True supports block streaming.

• ops.filtfilt(b, a, x) — zero-phase forward-backward filtering built on lfilter. Thus it inherits the custom VJP. It matches scipy.signal.filtfilt’s defaults: odd-reflection padding and steady-state edge conditions.

• ops.state_space(A, B, C, D, x, zi) — the general MIMO linear state space y[n] = C[n]·h[n] + D[n]·x[n], h[n+1] = A[n]·h[n] + B[n]·x[n]. It follows philtorch’s convention: LTI when the maps are constant, LPV when the maps are time-varying. Its custom VJP is the control-theory adjoint: the transposed state space run backwards (λ[n] = Cᵀg_y[n] + Aᵀλ[n+1]). This rule generalizes allpole’s companion-form rule to any state dimension.

• ops.diag_state_space(a, B, C, D, x, zi) — the diagonal/modal form (philtorch’s diag_state_space). With A = diag(a), the recursion decouples into N independent one-pole modes. Thus it runs as N linear_recurrence calls: O(T·N) instead of the dense O(T·N²), with the same custom VJP. To make a real resonant system, pass complex eigenvalues in a with complex B/C. The imaginary part of the output cancels; take jnp.real.

• faustax.dynamics — torchcomp’s user-facing gain functions built on these primitives: compexp_gain, limiter_gain, avg, ms2coef/coef2ms, amp2db/db2amp (unbatched; jax.vmap adds batching).

• faustax.filters — a differentiable TPT state-variable filter built on state_space. svf(x, cutoff, q, sample_rate, mode) is the functional core (lowpass/bandpass/highpass/notch). The cutoff and q are scalar or per-sample values for modulation. SVF is a learnable nnx.Module. Its cutoff and resonance train through smooth unconstrained parameters (from faustax import SVF, svf).

The forward outputs and all gradients agree with the torch implementations to machine precision (~1e-15 in float64). tests/test_ops.py pins them against naive-scan autodiff. The table shows CPU timings from examples/benchmarks/bench_ops_vs_torch.py (median ms). The benchmark compares torch 2.13 with torchlpc’s C++/OpenMP extension against jitted JAX. The benchmark requires pip install torch torchlpc torchcomp. Faustax itself does not depend on these packages.

case (1 s @ 44.1 kHz unless noted)	torch fwd	jax fwd	torch grad	jax grad
all-pole order 2, B=1	0.4	0.4	1.0	0.8
all-pole order 2, B=8	0.3	0.9	2.7	2.3
all-pole order 2, B=8, 10 s	3.6	10.3	22.5	26.1
ballistics, B=1	0.2	0.2	1.5	0.6
ballistics, B=8	0.3	1.1	2.8	2.2
compexp gain chain, B=1	—	—	2.7	0.8
compexp gain chain, B=8	—	—	4.6	3.2
biquad DF2 lfilter, B=1	—	—	1.5	2.2
biquad DF2 lfilter, B=8	—	—	4.8	6.1

At these sizes, the JAX gradients are equal in speed or faster. Torch is faster in the long-signal forward runs because of its per-batch OpenMP threading. The compexp rows run torchcomp’s full compexp_gain chain with identical math on the two sides. The biquad rows compare against philtorch’s compiled C++ lti.lfilter with unroll_factor=1, its fastest path. The default round(sqrt(T)) selects a pure-PyTorch fallback. This fallback measures ~3x slower than our JAX port. philtorch’s dedicated kernel stays ~1.4x faster than our scan composition. For comparison, BPTT through the Faust-generated Compressor costs 1.5 ms at B=1 but 59 ms at B=8. The Compressor is a different, richer topology: a switched peak detector and dB-domain smoothing. The 59 ms is ~19x more than the custom-VJP gain chain. Use the primitive when a training loss needs only compressor-style gain dynamics.

examples/benchmarks/bench_state_space.py (pure JAX, no torch) compares the state-space ops. The dense state_space adjoint runs at the same speed as plain BPTT-through-scan. XLA already differentiates a linear scan efficiently. Thus the benefit of the custom rule there is a bounded-memory residual and clean higher-order composition, not CPU speed. The modal diag_state_space gives the forward speedup. The decoupled modes cost O(T·N) and the dense form costs O(T·N²). We measured a speedup of ~17x (N=64) to ~63x (N=256) over 1 s @ 44.1 kHz. examples/parameter_estimation/fit_state_space_filter.py fits the readout of a state-space filter to target audio by gradient descent. It is a convex, exact-recovery demo.

GPU (uv sync --group gpu): sequential scans are latency-bound. Thus the CPU is faster below ~64 parallel lanes (EQ batch 1: 76 ms CPU vs 880 ms GPU). Above that count, the flat batch scaling of the GPU is faster. 1,024 vmapped DX7 synth voices cost the same wall time as one voice (461x real-time aggregate on an RTX 4080 SUPER). See the details and the DX7 study in the NNX backend notes.