Address review comments

Author: melissawm

docs/development/contribute_docs.md

@@ -26,15 +26,17 @@ If you are writing documentation for MaxText, you may want to preview the
 documentation site locally to ensure things work as expected before a deployment
 to [Read The Docs](https://readthedocs.org/).
 
-First, make sure you [install MaxText from source](install-from-source) and
-install the necessary dependencies. You can do this by navigating to your local
-clone of the MaxText repo and running:
+First, make sure you
+[install MaxText from source](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source)
+and install the necessary dependencies. You can do this by navigating to your
+local clone of the MaxText repo and running:
 
 ```bash
 uv pip install -r src/dependencies/requirements/requirements_docs.txt
 ```
 
-Once the dependencies are installed and your `maxtext_venv` virtual environment is activated, you can navigate to the `docs/` folder and run:
+Once the dependencies are installed and your `maxtext_venv` virtual environment
+is activated, you can navigate to the `docs/` folder and run:
 
 ```bash
 sphinx-build -b html . _build/html
@@ -61,6 +63,7 @@ toctree.
 
 ## Documentation deployment
 
-The MaxText documentation is deployed to
-[https://maxtext.readthedocs.io](https://maxtext.readthedocs.io) on any
-successful merge to the main branch.
+The latest version of the MaxText documentation, tracking the main branch of
+development, is automatically deployed to
+[https://maxtext.readthedocs.io/en/latest](https://maxtext.readthedocs.io/en/latest)
+on any successful merge to the main branch.

docs/install_maxtext.md

@@ -19,7 +19,8 @@
 This document discusses how to install MaxText.
 
 We recommend installing MaxText inside a Python virtual environment and using
-the [uv package manager](https://docs.astral.sh/uv/getting-started/installation/).
+the `uv` package manager following
+[uv's official installation instructions](https://docs.astral.sh/uv/getting-started/installation/).
 
 ```{note}
 MaxText is only tested on Linux during releases.
@@ -29,25 +30,16 @@ MaxText is only tested on Linux during releases.
 
 This is the easiest way to get started with the latest stable version.
 
-1. **Clone the repository:**
+1. **Create a virtual environment:**
 
    ```bash
-   git clone https://github.com/AI-Hypercomputer/maxtext.git
-   cd maxtext
+   uv venv --python 3.12 --seed <virtual env name>
+   source <virtual env name>/bin/activate
    ```
 
-2. **Create a virtual environment:**
-
-   ```bash
-   uv venv --python 3.12 --seed maxtext_venv
-   source maxtext_venv/bin/activate
-   ```
+2. **Install MaxText and its dependencies.**
 
-3. **Install MaxText and its dependencies.**
-
-   Choose a single installation option from this list to fit your use case. For
-   each installation option, a different helper script is necessary to complete
-   the installation of required dependencies that are not yet available on PyPI.
+   Choose a single installation option from this list to fit your use case.
 
    ```{important}
    If you want to switch to a different installation option (e.g., from `[tpu]`
@@ -60,30 +52,30 @@ This is the easiest way to get started with the latest stable version.
 
      ```bash
      uv pip install maxtext[tpu]==0.2.1 --resolution=lowest
-     install_tpu_pre_train_extra_deps
      ```
 
    - **Option 2:** Install `maxtext[cuda12]`, used for pre-training and decoding
      on GPUs.
 
      ```bash
      uv pip install maxtext[cuda12]==0.2.1 --resolution=lowest
-     install_cuda12_pre_train_extra_deps
      ```
 
    - **Option 3:** Install `maxtext[tpu-post-train]`, used for post-training on
-     TPUs. Currently, this option should also be used for running `vllm_decode` on TPUs.
+     TPUs. Currently, this option should also be used for running `vllm_decode`
+     on TPUs.
 
      ```bash
      uv pip install maxtext[tpu-post-train]==0.2.1 --resolution=lowest
-     install_tpu_post_train_extra_deps
      ```
 
    - **Option 4:** Install `maxtext[runner]`, used for building MaxText's Docker
      images and scheduling workloads through XPK. Once installed, you will have
      access to the `build_maxtext_docker_image`, `upload_maxtext_docker_image`,
      and `xpk` commands. For more details on building and uploading Docker
-     images, see the [Build MaxText Docker Image](https://maxtext.readthedocs.io/en/latest/build_maxtext.html) guide.
+     images, see the
+     [Build MaxText Docker Image](https://maxtext.readthedocs.io/en/latest/build_maxtext.html)
+     guide.
 
      ```bash
      uv pip install maxtext[runner]==0.2.1 --resolution=lowest
@@ -123,8 +115,8 @@ environment to avoid dependency conflicts.
 2. Create virtual environment:
 
    ```bash
-   uv venv --python 3.12 --seed maxtext_venv
-   source maxtext_venv/bin/activate
+   uv venv --python 3.12 --seed <virtual env name>
+   source <virtual env name>/bin/activate
    ```
 
 3. Install dependencies in editable mode. Choose a single installation option
@@ -157,4 +149,6 @@ environment to avoid dependency conflicts.
      uv pip install -e .[runner] --resolution=lowest
      ```
 
-After installation, you can verify the package is available with `python3 -c "import maxtext"` and run training jobs with `python3 -m maxtext.trainers.pre_train.train ...`.
+After installation, you can verify the package is available with
+`python3 -c "import maxtext"` and run training jobs with
+`python3 -m maxtext.trainers.pre_train.train ...`.