jax.Array.swapaxes