<!--
 Copyright 2023-2025 Google LLC

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

      https://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 -->

# Install MaxText

This document discusses how to install MaxText. We recommend installing MaxText inside a Python virtual environment.

## From PyPI (Recommended)
This is the easiest way to get started with the latest stable version.

```bash
# 1. Install uv, a fast Python package installer
pip install uv

# 2. Create virtual environment
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

# 3. Install MaxText and its dependencies
uv pip install maxtext --resolution=lowest
install_maxtext_github_deps
```
> **Note:** The `install_maxtext_github_deps` command is temporarily required to install dependencies directly from GitHub that are not yet available on PyPI.

> **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
# 1. Clone the repository
git clone https://github.com/AI-Hypercomputer/maxtext.git
cd maxtext

# 2. Create virtual environment
pip install uv
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

# 3. Install dependencies in editable mode
# install the tpu package
uv pip install -e .[tpu] --resolution=lowest
# or install the gpu package by running the following line
# uv pip install -e .[cuda12] --resolution=lowest
install_maxtext_github_deps
```

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 ...`.

# Update MaxText dependencies

## Introduction

This document provides a guide to updating dependencies in MaxText using the `seed-env` tool. `seed-env` helps generate deterministic and reproducible Python environments by creating fully-pinned `requirements.txt` files from a base set of requirements.

## Overview of the Process

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.

The following sections provide detailed instructions for each step.

## Step 1: Install seed-env

First, you need to install the `seed-env` command-line tool by running `pip install seed-env uv`. Or follow the instructions in the
[seed-env repository](https://github.com/google-ml-infra/actions/tree/main/python_seed_env#install-the-seed-env-tool) if you want to build `seed-env` from source.

## Step 2: Find the JAX Build Commit Hash

The dependency generation process is pinned to a specific nightly build of JAX. You need to find the commit hash for the desired JAX build.

You can find the latest commit hashes in the [JAX `build/` folder](https://github.com/jax-ml/jax/commits/main/build). Choose a recent, successful build and copy its full commit hash.

## Step 3: Generate the Requirements Files

Next, run the `seed-env` CLI to generate the new requirements files. You will need to do this separately for the TPU and GPU environments. The generated files will be placed in a directory specified by `--output-dir`.

### For TPU

Run the following command, replacing `<jax-build-commit-hash>` with the hash you copied in the previous step.

```bash
seed-env \
  --local-requirements=dependencies/requirements/base_requirements/tpu-base-requirements.txt \
  --host-name=MaxText \
  --seed-commit=<jax-build-commit-hash> \
  --python-version=3.12 \
  --requirements-txt=tpu-requirements.txt \
  --output-dir=generated_tpu_artifacts
```

### For GPU

Similarly, run the command for the GPU requirements.

```bash
seed-env \
  --local-requirements=dependencies/requirements/base_requirements/cuda12-base-requirements.txt \
  --host-name=MaxText \
  --seed-commit=<jax-build-commit-hash> \
  --python-version=3.12 \
  --requirements-txt=cuda12-requirements.txt \
  --hardware=cuda12 \
  --output-dir=generated_gpu_artifacts
```

## Step 4: Update Project Files

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`.

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.

```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.

```bash
pip install uv
# install the tpu package
uv pip install -e .[tpu] --resolution=lowest
# or install the gpu package by running the following line:
# uv pip install -e .[cuda12] --resolution=lowest
install_maxtext_github_deps
```

3.  **Run tests:** Run MaxText tests to ensure there are no regressions.