[JAX] Integrate Permutation to MaxText #2585
Description
After permutation primitive is merged, needs to integrate to MaxText for full model testing
{{ message }}
After permutation primitive is merged, needs to integrate to MaxText for full model testing