Specifically install maxtext 0.2.1 in documentation

[bvandermoon]

Description

Pin maxtext==0.2.1 in ReadTheDocs installation guides. Our commands have been defaulting to install maxtext==0.1.0, which we figured out is caused by using the --resolution=lowest flag to keep transitive dependencies locked.

Pinning the latest version in the docs for now, but we will want to find a better long-term solution.

Next steps: Can we somehow cherry-pick this change into the ReadTheDocs 0.2.1 tag?

Tests

uv pip install maxtext[tpu]==0.2.1 --resolution=lowest
pip show maxtext # gives 0.2.1

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.