jax.lax.fft =========== .. currentmodule:: jax.lax .. autofunction:: fft