jax

8377224b - ppermute (aka collective_permute) jvp + transpose