Fix MaxText installation instructions in all jupyter notebooks and update doc

.github/workflows/run_jupyter_notebooks.yml

@@ -64,8 +64,6 @@ jobs:
 
           # 2. Install MaxText package and all the post training dependencies
           uv pip install ${maxtext_wheel}[tpu-post-train] --resolution=lowest
-          #TODO: @mazumdera:  replace this with the following after release
-          # uv pip install maxtext[tpu-post-train] --resolution=lowest
           install_maxtext_tpu_post_train_extra_deps
           .venv/bin/python3 -m ipykernel install --user --name maxtext_venv
 

docs/guides/run_python_notebook.md

@@ -103,7 +103,7 @@ To install, click the `Extensions` icon on the left sidebar (or press `Ctrl+Shif
 
 ### Step 4: Install MaxText and Dependencies
 
-To execute post-training notebooks on your TPU-VM, follow the official [MaxText installation guides](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl.html#create-virtual-environment-and-install-maxtext-dependencies) to install MaxText and its dependencies inside a dedicated virtual environment.
+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.
 
 ### Step 5: Install the necessary library for Jupyter
 
@@ -162,7 +162,7 @@ pip3 install jupyterlab
 
 ### Step 4: Install MaxText and Dependencies
 
-To execute post-training notebooks on your TPU-VM, follow the official [MaxText installation guides](https://maxtext.readthedocs.io/en/latest/tutorials/posttraining/rl.html#create-virtual-environment-and-install-maxtext-dependencies) to install MaxText and its dependencies inside a dedicated virtual environment.
+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.
 
 ### Step 5: Register virtual environment as a Jupyter Kernel
 

docs/install_maxtext.md

@@ -24,6 +24,7 @@ MaxText offers three installation modes:
 3. maxtext[tpu-post-train]. Used for post-training on TPUs. Currently, this option should also be used for running vllm_decode on TPUs.
 
 ## From PyPI (Recommended)
+
 This is the easiest way to get started with the latest stable version.
 
 ```bash
@@ -38,24 +39,26 @@ source maxtext_venv/bin/activate
 #      installation option from this list to fit your use case.
 
 # Option 1: Installing maxtext[tpu]
-uv pip install maxtext[tpu] --resolution=lowest
+uv pip install "maxtext[tpu]>=0.2.0" --resolution=lowest
 install_maxtext_tpu_github_deps
 
 # Option 2: Installing maxtext[cuda12]
-uv pip install maxtext[cuda12] --resolution=lowest
+uv pip install "maxtext[cuda12]>=0.2.0" --resolution=lowest
 install_maxtext_cuda12_github_dep
 
 # Option 3: Installing maxtext[tpu-post-train]
-uv pip install maxtext[tpu-post-train] --resolution=lowest
+uv pip install "maxtext[tpu-post-train]>=0.2.0" --resolution=lowest
 install_maxtext_tpu_post_train_extra_deps
 ```
+
 > **Note:** The `install_maxtext_tpu_github_deps`, `install_maxtext_cuda12_github_dep`, and
-`install_maxtext_tpu_post_train_extra_deps` commands are temporarily required to install dependencies directly from GitHub
-that are not yet available on PyPI. As shown above, choose the one that corresponds to your use case.
+> `install_maxtext_tpu_post_train_extra_deps` commands are temporarily required to install dependencies directly from GitHub
+> that are not yet available on PyPI. As shown above, choose the one that corresponds to your use case.
 
 > **Note:** The maxtext package contains a comprehensive list of all direct and transitive dependencies, with lower bounds, generated by [seed-env](https://github.com/google-ml-infra/actions/tree/main/python_seed_env). We highly recommend the `--resolution=lowest` flag. It instructs `uv` to install the specific, tested versions of dependencies defined by MaxText, rather than the latest available ones. This ensures a consistent and reproducible environment, which is critical for stable performance and for running benchmarks.
 
 ## From Source
+
 If you plan to contribute to MaxText or need the latest unreleased features, install from source.
 
 ```bash
@@ -98,11 +101,11 @@ Please keep dependencies updated throughout development. This will allow each co
 
 To update dependencies, you will follow these general steps:
 
-1.  **Modify Base Requirements**: Update the desired dependencies in `base_requirements/requirements.txt` or the hardware-specific files (`base_requirements/tpu-base-requirements.txt`, `base_requirements/gpu-base-requirements.txt`).
-2.  **Generate New Files**: Run the `seed-env` CLI tool to generate new, fully-pinned requirements files based on your changes.
-3.  **Update Project Files**: Copy the newly generated files into the `generated_requirements/` directory.
-4.  **Handle GitHub Dependencies**: Move any dependencies that are installed directly from GitHub from the generated files to `src/install_maxtext_extra_deps/extra_deps_from_github.txt`.
-5.  **Verify**: Test the new dependencies to ensure the project installs and runs correctly.
+1. **Modify Base Requirements**: Update the desired dependencies in `base_requirements/requirements.txt` or the hardware-specific files (`base_requirements/tpu-base-requirements.txt`, `base_requirements/gpu-base-requirements.txt`).
+2. **Generate New Files**: Run the `seed-env` CLI tool to generate new, fully-pinned requirements files based on your changes.
+3. **Update Project Files**: Copy the newly generated files into the `generated_requirements/` directory.
+4. **Handle GitHub Dependencies**: Move any dependencies that are installed directly from GitHub from the generated files to `src/install_maxtext_extra_deps/extra_deps_from_github.txt`.
+5. **Verify**: Test the new dependencies to ensure the project installs and runs correctly.
 
 The following sections provide detailed instructions for each step.
 
@@ -154,25 +157,26 @@ seed-env \
 
 After generating the new requirements, you need to update the files in the MaxText repository.
 
-1.  **Copy the generated files:**
-    -   Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`.
-    -   Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`.
+1. **Copy the generated files:**
+
+   - Move `generated_tpu_artifacts/tpu-requirements.txt` to `generated_requirements/tpu-requirements.txt`.
+   - Move `generated_gpu_artifacts/cuda12-requirements.txt` to `generated_requirements/cuda12-requirements.txt`.
 
-2.  **Update `extra_deps_from_github.txt` (if necessary):**
-    Currently, MaxText uses a few dependencies, such as `mlperf-logging` and `google-jetstream`, that are installed directly from GitHub source. These are defined in `base_requirements/requirements.txt`, and the `seed-env` tool will carry them over to the generated requirements files.
+2. **Update `extra_deps_from_github.txt` (if necessary):**
+   Currently, MaxText uses a few dependencies, such as `mlperf-logging` and `google-jetstream`, that are installed directly from GitHub source. These are defined in `base_requirements/requirements.txt`, and the `seed-env` tool will carry them over to the generated requirements files.
 
 ## Step 5: Verify the New Dependencies
 
 Finally, test that the new dependencies install correctly and that MaxText runs as expected.
 
-1.  **Create a clean environment:** It's best to start with a fresh Python virtual environment.
+1. **Create a clean environment:** It's best to start with a fresh Python virtual environment.
 
 ```bash
 uv venv --python 3.12 --seed maxtext_venv
 source maxtext_venv/bin/activate
 ```
 
-2.  **Run the setup script:** Execute `bash setup.sh` to install the new dependencies.
+2. **Run the setup script:** Execute `bash setup.sh` to install the new dependencies.
 
 ```bash
 pip install uv
@@ -183,4 +187,4 @@ uv pip install -e .[tpu] --resolution=lowest
 install_maxtext_github_deps
 ```
 
-3.  **Run tests:** Run MaxText tests to ensure there are no regressions.
+3. **Run tests:** Run MaxText tests to ensure there are no regressions.

docs/tutorials/posttraining/rl.md

@@ -57,7 +57,7 @@ source ${VENV_NAME?}/bin/activate
 Run the following commands to get all the necessary installations.
 
 ```bash
-uv pip install maxtext[tpu-post-train] --resolution=lowest
+uv pip install "maxtext[tpu-post-train]>=0.2.0" --resolution=lowest
 install_maxtext_tpu_post_train_extra_deps
 ```
 

docs/tutorials/posttraining/sft.md

@@ -37,7 +37,7 @@ source ${VENV_NAME?}/bin/activate
 Run the following commands to get all the necessary installations.
 
 ```bash
-uv pip install maxtext[tpu-post-train] --resolution=lowest
+uv pip install "maxtext[tpu-post-train]>=0.2.0" --resolution=lowest
 install_maxtext_tpu_post_train_extra_deps
 ```
 

src/maxtext/examples/rl_llama3_demo.ipynb

@@ -89,17 +89,38 @@
    "source": [
     "## Installation: MaxText and Post training Dependencies\n",
     "\n",
-    "**Running the notebook on Visual Studio or JupyterLab:** Create an virtual environment and install dependencies outside of the notebook using the commands in [MaxText installation and dependency setup guide](../../../docs/guides/run_python_notebook.md#step-4-install-maxtext-and-dependencies) before proceeding. And then run the notebook using that virtual environment.\n",
-    "\n",
-    "\n",
-    "**Running the notebook on Google Colab:** Run the above-mentioned commands in the Colab and then restart session\n",
+    "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if IN_COLAB:\n",
+    "    # Clone the MaxText repository\n",
+    "    !git clone https://github.com/AI-Hypercomputer/maxtext.git\n",
+    "    %cd maxtext\n",
     "\n",
+    "    # Install uv, a fast Python package installer\n",
+    "    !pip install uv\n",
+    "    \n",
+    "    # Install MaxText and post-training dependencies\n",
+    "    !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
+    "    !install_maxtext_tpu_post_train_extra_deps"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
     "**Session restart Instructions for Colab:**\n",
     "1.  Navigate to the menu at the top of the screen.\n",
     "2.  Click on **Runtime**.\n",
     "3.  Select **Restart session** from the dropdown menu.\n",
     "\n",
-    "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**.\n"
+    "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
    ]
   },
   {
@@ -119,20 +140,19 @@
     "import os\n",
     "import sys\n",
     "from pathlib import Path\n",
-    "import MaxText\n",
     "from huggingface_hub import login\n",
     "from etils import epath\n",
     "import jax\n",
     "\n",
     "from maxtext.trainers.post_train.rl.train_rl import rl_train, setup_configs_and_devices\n",
+    "from maxtext.utils.globals import MAXTEXT_REPO_ROOT, MAXTEXT_PKG_DIR\n",
     "\n",
     "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"0\"\n",
     "os.environ[\"SKIP_JAX_PRECOMPILE\"] = \"1\"  # Faster startup for vLLM\n",
     "# Suppress vLLM logging with a severity level below ERROR\n",
     "os.environ[\"VLLM_LOGGING_LEVEL\"] = \"ERROR\"\n",
     "\n",
-    "MAXTEXT_PKG_DIR = os.path.dirname(MaxText.__file__)\n",
-    "MAXTEXT_REPO_ROOT = os.sep.join([\"maxtext\" if p == \"MaxText\" else p for p in MAXTEXT_PKG_DIR.split(os.sep)])\n",
+    "\n",
     "print(f\"MaxText installation path: {MAXTEXT_PKG_DIR}\")"
    ]
   },
@@ -154,9 +174,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "if IN_COLAB:\n",
+    "try:\n",
+    "    from google.colab import userdata\n",
     "    HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
-    "else:\n",
+    "except ImportError:\n",
     "    HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
     "\n",
     "# If not found in the environment, prompt the user for input securely\n",
@@ -191,7 +212,7 @@
     "RUN_NAME = datetime.datetime.now().strftime(\"%Y-%m-%d-%H-%M-%S\")\n",
     "LOSS_ALGO=\"grpo\" #  or \"gspo-token\" if you want to use GSPO\n",
     "\n",
-    "CHAT_TEMPLATE_PATH = f\"{MAXTEXT_REPO_ROOT}/examples/chat_templates/gsm8k_rl.json\"\n",
+    "CHAT_TEMPLATE_PATH = f\"{MAXTEXT_REPO_ROOT}/src/maxtext/examples/chat_templates/gsm8k_rl.json\"\n",
     "if not os.path.exists(CHAT_TEMPLATE_PATH):\n",
     "    raise FileNotFoundError(f\"Chat template not found: {CHAT_TEMPLATE_PATH}\")\n",
     "\n",

src/maxtext/examples/sft_llama3_demo_tpu.ipynb

@@ -101,17 +101,38 @@
       "source": [
         "## Installation: MaxText and Post training Dependencies\n",
         "\n",
-        "**Running the notebook on Visual Studio or JupyterLab:** Create an virtual environment and install dependencies outside of the notebook using the commands in [MaxText installation and dependency setup guide](../../../docs/guides/run_python_notebook.md#step-4-install-maxtext-and-dependencies) before proceeding. And then run the notebook using that virtual environment.\n",
-        "\n",
-        "\n",
-        "**Running the notebook on Google Colab:** Run the above-mentioned commands in the Colab and then restart session\n",
-        "\n",
+        "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "if IN_COLAB:\n",
+        "    # Clone the MaxText repository\n",
+        "    !git clone https://github.com/AI-Hypercomputer/maxtext.git\n",
+        "    %cd maxtext\n",
+        "\n",
+        "    # Install uv, a fast Python package installer\n",
+        "    !pip install uv\n",
+        "    \n",
+        "    # Install MaxText and post-training dependencies\n",
+        "    !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
+        "    !install_maxtext_tpu_post_train_extra_deps"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
         "**Session restart Instructions for Colab:**\n",
         "1.  Navigate to the menu at the top of the screen.\n",
         "2.  Click on **Runtime**.\n",
         "3.  Select **Restart session** from the dropdown menu.\n",
         "\n",
-        "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**.\n"
+        "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
       ]
     },
     {
@@ -166,9 +187,10 @@
       },
       "outputs": [],
       "source": [
-        "if IN_COLAB:\n",
+        "try:\n",
+        "    from google.colab import userdata\n",
         "    HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
-        "else:\n",
+        "except ImportError:\n",
         "    HF_TOKEN = os.environ.get(\"HF_TOKEN\", \"\")\n",
         "\n",
         "# If not found in the environment, prompt the user for input securely\n",
@@ -250,7 +272,6 @@
         "    print(\"Converting checkpoint from HuggingFace...\")\n",
         "    env = os.environ.copy()\n",
         "    env[\"JAX_PLATFORMS\"] = \"cpu\"\n",
-        "    # env[\"PYTHONPATH\"] = MAXTEXT_PKG_DIR\n",
         "\n",
         "    subprocess.run(\n",
         "        [\n",

src/maxtext/examples/sft_qwen3_demo.ipynb

@@ -95,40 +95,46 @@
     "    IN_COLAB = False"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "D9ms-jTSZUQL"
+   },
+   "source": [
+    "## Installation: MaxText and Post training Dependencies\n",
+    "\n",
+    "**Running the notebook on Visual Studio or JupyterLab:** Before proceeding, create a virtual environment and install the required post-training dependencies by following `Option 3: Installing [tpu-post-train]` in the [MaxText installation guide](https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-source). Once the environment is set up, ensure the notebook is running within it."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "try:\n",
-    "    import google.colab\n",
-    "    print(\"Running the notebook in Google Colab\")\n",
-    "    IN_COLAB = True\n",
-    "except ImportError:\n",
-    "    print(\"Running the notebook on JupyterLab\")\n",
-    "    IN_COLAB = False"
+    "if IN_COLAB:\n",
+    "    # Clone the MaxText repository\n",
+    "    !git clone https://github.com/AI-Hypercomputer/maxtext.git\n",
+    "    %cd maxtext\n",
+    "\n",
+    "    # Install uv, a fast Python package installer\n",
+    "    !pip install uv\n",
+    "    \n",
+    "    # Install MaxText and post-training dependencies\n",
+    "    !uv pip install -e .[tpu-post-train] --resolution=lowest\n",
+    "    !install_maxtext_tpu_post_train_extra_deps"
    ]
   },
   {
    "cell_type": "markdown",
-   "metadata": {
-    "id": "D9ms-jTSZUQL"
-   },
+   "metadata": {},
    "source": [
-    "## Installation: MaxText and Post training Dependencies\n",
-    "\n",
-    "**Running the notebook on Visual Studio or JupyterLab:** Create an virtual environment and install dependencies outside of the notebook using the commands in [MaxText installation and dependency setup guide](../../../docs/guides/run_python_notebook.md#step-4-install-maxtext-and-dependencies) before proceeding. And then run the notebook using that virtual environment.\n",
-    "\n",
-    "\n",
-    "**Running the notebook on Google Colab:** Run the above-mentioned commands in the Colab and then restart session\n",
-    "\n",
     "**Session restart Instructions for Colab:**\n",
     "1.  Navigate to the menu at the top of the screen.\n",
     "2.  Click on **Runtime**.\n",
     "3.  Select **Restart session** from the dropdown menu.\n",
     "\n",
-    "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**.\n"
+    "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
    ]
   },
   {
@@ -153,7 +159,6 @@
     "import sys\n",
     "import transformers\n",
     "\n",
-    "import MaxText\n",
     "from maxtext.configs import pyconfig\n",
     "from maxtext.examples.sft_train_and_evaluate import evaluate_model, get_test_dataset\n",
     "from maxtext.integration.tunix.tunix_adapter import TunixMaxTextAdapter\n",

src/maxtext/examples/sft_train_and_evaluate.py

@@ -30,7 +30,7 @@
 
 # Run the following commands to get all the necessary installations.
 
-uv pip install maxtext[tpu-post-train] --resolution=lowest
+uv pip install "maxtext[tpu-post-train]>=0.2.0" --resolution=lowest
 install_maxtext_tpu_post_train_extra_deps