Unable to install in Mac/M2

[melissawm]

Bug report

While following the instructions in the "Install MaxText" page on an M2/arm Mac, I get a dependency resolution error.

Installing with only pip install -e . works, but this option is not listed in the installation instructions. I think this is worth mentioning, especially as building the docs will require installing MaxText locally, but does not require a TPU or GPU.

Logs/Output

~/projects/maxtext (docgen ✗) uv venv --python 3.12 --seed maxtext_venv
Using CPython 3.12.11 interpreter at: /opt/homebrew/opt/python@3.12/bin/python3.12
Creating virtual environment with seed packages at: maxtext_venv
 + pip==25.3
Activate with: source maxtext_venv/bin/activate
~/projects/maxtext (docgen ✗) source maxtext_venv/bin/activate
(maxtext_venv) ~/projects/maxtext (docgen ✗) uv pip install -e ."[tpu]" --resolution=lowest
Using Python 3.12.11 environment at: maxtext_venv
      Built maxtext @ file:///Users/mwebermendonca/projects/maxtext
  × No solution found when resolving dependencies:
  ╰─▶ Because only the following versions of triton are available:
          triton<=3.5.0
          triton==3.5.1
      and triton>=3.5.0 has no wheels with a matching platform tag (e.g., `macosx_26_0_arm64`), we can conclude that triton>=3.5.0 cannot be used.
      And because maxtext[tpu]==0.1.1 depends on triton>=3.5.0, we can conclude that maxtext[tpu]==0.1.1 cannot be used.
      And because only maxtext[tpu]==0.1.1 is available and you require maxtext[tpu], we can conclude that your requirements are unsatisfiable.

      hint: Wheels are available for `triton` (v3.5.0) on the following platforms: `manylinux_2_27_aarch64`, `manylinux_2_27_x86_64`, `manylinux_2_28_aarch64`, `manylinux_2_28_x86_64`
(maxtext_venv) ~/projects/maxtext (docgen ✗) 

Environment Information

~ sysctl machdep.cpu
machdep.cpu.cores_per_package: 8
machdep.cpu.core_count: 8
machdep.cpu.logical_per_package: 8
machdep.cpu.thread_count: 8
machdep.cpu.brand_string: Apple M2

Additional Context

No response