Merge pull request #3594 from AI-Hypercomputer:bvandermoon-uxr-fixes

PiperOrigin-RevId: 896127826

Author: Google-ML-Automation

docs/guides/run_python_notebook.md

@@ -105,6 +105,8 @@ To install, click the `Extensions` icon on the left sidebar (or press `Ctrl+Shif
 
 To execute post-training notebooks on your TPU-VM, follow the official [MaxText installation guides](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source) and specifically follow `Option 3: Installing [tpu-post-train]`. This will ensure all post-training dependencies are installed inside your virtual environment.
 
+> **Note:** If you have previously installed MaxText with a different option (e.g., `maxtext[tpu]`), we strongly recommend using a fresh virtual environment for `maxtext[tpu-post-train]` to avoid potential library version conflicts.
+
 ### Step 5: Install the necessary library for Jupyter
 
 Jupyter requires a kernel to execute code. This kernel is tied to a specific Python environment. Open your terminal inside VS Code and run:
@@ -164,6 +166,8 @@ pip3 install jupyterlab
 
 To execute post-training notebooks on your TPU-VM, follow the official [MaxText installation guides](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source) and specifically follow `Option 3: Installing [tpu-post-train]`. This will ensure all post-training dependencies are installed inside your virtual environment.
 
+> **Note:** If you have previously installed MaxText with a different option (e.g., `maxtext[tpu]`), we strongly recommend using a fresh virtual environment for `maxtext[tpu-post-train]` to avoid potential library version conflicts.
+
 ### Step 5: Register virtual environment as a Jupyter Kernel
 
 Once the environment is set up, you need to register it so JupyterLab can see it as an available kernel. Run the following command on your TPU-VM, replacing <virtual env name> with the actual name to your virtual environment:

docs/install_maxtext.md

@@ -38,6 +38,9 @@ source maxtext_venv/bin/activate
 
 # 3. Install MaxText and its dependencies. 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] to [tpu-post-train]), we strongly recommend
+# starting with a fresh virtual environment to avoid dependency conflicts.
 
 # Option 1: Installing maxtext[tpu]
 uv pip install maxtext[tpu]==0.2.1 --resolution=lowest
@@ -80,7 +83,10 @@ uv venv --python 3.12 --seed maxtext_venv
 source maxtext_venv/bin/activate
 
 # 3. Install dependencies in editable mode. Choose a single
-#      installation option from this list to fit your use case.
+# 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] to [tpu-post-train]), we strongly recommend
+# starting with a fresh virtual environment to avoid dependency conflicts.
 
 # Option 1: Installing .[tpu]
 uv pip install -e .[tpu] --resolution=lowest

docs/tutorials/posttraining/rl.md

@@ -46,6 +46,8 @@ Let's get started!
 
 For instructions on installing MaxText with post-training dependencies on your VM, please refer to the [official documentation](https://maxtext.readthedocs.io/en/latest/install_maxtext.html) and use the `maxtext[tpu-post-train]` installation path to include all necessary post-training dependencies.
 
+> **Note:** If you have previously installed MaxText with a different option (e.g., `maxtext[tpu]`), we strongly recommend using a fresh virtual environment for `maxtext[tpu-post-train]` to avoid potential library version conflicts.
+
 ## Setup environment variables
 
 Login to Hugging Face. Provide your access token when prompted:

docs/tutorials/posttraining/sft.md

@@ -28,6 +28,8 @@ In this tutorial we use a single host TPU VM such as `v6e-8/v5p-8`. Let's get st
 
 For instructions on installing MaxText with post-training dependencies on your VM, please refer to the [official documentation](https://maxtext.readthedocs.io/en/latest/install_maxtext.html) and use the `maxtext[tpu-post-train]` installation path to include all necessary post-training dependencies.
 
+> **Note:** If you have previously installed MaxText with a different option (e.g., `maxtext[tpu]`), we strongly recommend using a fresh virtual environment for `maxtext[tpu-post-train]` to avoid potential library version conflicts.
+
 ## Setup environment variables
 
 Login to Hugging Face. Provide your access token when prompted: