[JAX] Integrate Permutation to MaxText

[tdophung]

After permutation primitive is merged, needs to integrate to MaxText for full model testing