jax.lax.all\_to\_all ==================== .. currentmodule:: jax.lax .. autofunction:: all_to_all