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

[festusev]

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

[melissawm]

It also looks like the setup instructions in this document do not match the ones in https://maxtext.readthedocs.io/en/latest/guides/install_maxtext.html - we should probably direct everyone to follow the latter, now that this page exists.

Additionally, while following the instructions in the "Install MaxText" page, I have a different error but I'll open a separate issue for it.