jax.lax.pswapaxes ================= .. currentmodule:: jax.lax .. autofunction:: pswapaxes