Dependency conflicts between tpu-commons, tokamax, and google-tunix

### Bug report

When building from source following the commands in https://maxtext.readthedocs.io/en/latest/tutorials/grpo.html I get the following dependency conflict:
tpu-commons 0.1.2 requires jax==0.7.2 and jaxlib==0.7.2, however tokamax 0.0.4 requires jax and jaxlib >= 0.8.0.

Also, google-tunix==0.1.3 requires jax==0.8

### Logs/Output

_No response_

### Environment Information

The commands I have run are:
`python3 -m venv ~/venv-maxtext`
`source ~/venv-maxtext/bin/activate`
`bash tools/setup/setup.sh`
`pre-commit install`
`bash ~/maxtext/src/MaxText/examples/install_tunix_vllm_requirement.sh`

### Additional Context

_No response_