MaxText seems broken, as importing it throws and error from MaxText itself.

[syrkis]

Bug report

Just installed MaxText. Getting the error below. I'm on the newest version of MaxText (0.2.0)

Logs/Output

…!?] is 󰏗 v0.1.0 via  v3.12.8 on ☁️
❯ python main.py
Traceback (most recent call last):
  File "/Users/nobr/desk/nebellum/parabellum/main.py", line 14, in <module>
    import MaxText as mt
  File "/Users/nobr/desk/nebellum/parabellum/.venv/lib/python3.12/site-packages/MaxText/__init__.py", line 38, in <module>
    from maxtext.configs import pyconfig
ModuleNotFoundError: No module named 'maxtext'

Environment Information

MacOS
Python3.12

Additional Context

No response

[syrkis]

Seems to be associated with the maxtext restructuring. Is 0.2.0 considered broken, how am I going about this wrong?

[truluri]

Also having this issue

[bvandermoon]

Hey @syrkis @truluri, thank you for reporting this issue. I am able to import maxtext.configs with a clean maxtext[tpu]==0.2.0 install:

python3
Python 3.12.11 (main, Sep 18 2025, 19:47:19) [Clang 20.1.4 ] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import maxtext.configs

Could you please double-check this is running 0.2.0 with pip show maxtext? I see something in the original error message saying v0.1.0. We have seen 0.1.0 showing up with a standard pip install maxtext[tpu]

[syrkis]

Hi @bvandermoon. Thanks for replying. I now made a new venv, installed maxtext 0.2.0 and jax.
v0.1.0 in the original message is from uv init for the project I am doing (unrelated to maxtext). It was indeed on 0.2.0. Here is a print out from my shell, with me making a clean env, installing maxtext (and then jax) and then reproducing the issue.

desk/fun/maxtext_test on ☁️
❯ uv init --python 3.12
Initialized project `maxtext-test`

maxtext_test on  master [?] is 󰏗 v0.1.0 via  v3.14.3 on ☁️
❯ cat .python-version
─────┬────────────────────────────────────────────────────────────────────────────────────────────────────
     │ File: .python-version
─────┼────────────────────────────────────────────────────────────────────────────────────────────────────
   1 │ 3.12
─────┴────────────────────────────────────────────────────────────────────────────────────────────────────

maxtext_test on  master [?] is 󰏗 v0.1.0 via  v3.14.3 on ☁️
❯ uv add maxtext
Using CPython 3.12.8 interpreter at: /usr/local/bin/python3.12
Creating virtual environment at: .venv
Resolved 2 packages in 179ms
Installed 1 package in 29ms
 + maxtext==0.2.0

maxtext_test on  master [?] is 󰏗 v0.1.0 via  v3.14.3 on ☁️
❯ uv run python
Python 3.12.8 (v3.12.8:2dc476bcb91, Dec  3 2024, 14:43:19) [Clang 13.0.0 (clang-1300.0.29.30)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import maxtext.configs
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'maxtext'
>>> import MaxText
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/nobr/desk/fun/maxtext_test/.venv/lib/python3.12/site-packages/MaxText/__init__.py", line 36, in <module>
    from jax.sharding import Mesh
ModuleNotFoundError: No module named 'jax'
>>> ^D

maxtext_test on  master [?] is 󰏗 v0.1.0 via  v3.14.3 on ☁️   took 19s
❯ uv add maxtext --resolution=lowest
Ignoring existing lockfile due to change in resolution mode: `highest` vs. `lowest`
Resolved 2 packages in 28ms
Checked 1 package in 1ms

maxtext_test on  master [?] is 󰏗 v0.1.0 via  v3.14.3 on ☁️
❯ uv add jax
Ignoring existing lockfile due to change in resolution mode: `lowest` vs. `highest`
Resolved 8 packages in 187ms
Installed 6 packages in 126ms
 + jax==0.9.2
 + jaxlib==0.9.2
 + ml-dtypes==0.5.4
 + numpy==2.4.3
 + opt-einsum==3.4.0
 + scipy==1.17.1

maxtext_test on  master [?] is 󰏗 v0.1.0 via  v3.14.3 on ☁️
❯ uv run python
Python 3.12.8 (v3.12.8:2dc476bcb91, Dec  3 2024, 14:43:19) [Clang 13.0.0 (clang-1300.0.29.30)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import MaxText
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/nobr/desk/fun/maxtext_test/.venv/lib/python3.12/site-packages/MaxText/__init__.py", line 38, in <module>
    from maxtext.configs import pyconfig
ModuleNotFoundError: No module named 'maxtext'
>>>  import maxtext.configs
  File "<stdin>", line 1
    import maxtext.configs
IndentationError: unexpected indent
>>> import maxtext.configs
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'maxtext'
>>> import maxtext.configs

❯ uv pip show maxtext
Name: maxtext
Version: 0.2.0
Location: /Users/nobr/desk/fun/maxtext_test/.venv/lib/python3.12/site-packages
Requires:
Required-by:

[igorts-git]

Hi @syrkis could you try installing MaxText using the following instructions: https://maxtext.readthedocs.io/en/latest/install_maxtext.html#from-pypi-recommended
Use the option that is suitable for your use-case: (1) TPU pre-training (2) GPU (3) TPU Post-training.

I see that you use "import MaxText". We recently moved some code around and now it is supposed to be lowercase "import maxtext". Though, we tried to keep the old CamelCase import working as well. Could you try the lowercase?

[igorts-git]

@syrkis to follow up I tried using uv project to see if that works. Some commands need to be different from our instructions, but it seems to work. Here are the commands that I ran:

$ mkdir clean_install
$ cd clean_install
$ uv init
.. initialized the project ..
$ uv add "maxtext[tpu]>=0.2.0"
.. installed a bunch of packages ..
$ uv run python3
Python 3.12.12 (main, Jan 27 2026, 23:58:14) [Clang 21.1.4 ] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import MaxText
PyTorch was not found. Models won't be available and only tokenizers, configuration and file/data utilities can be used.
2026-03-19 11:56:50.265659: I external/local_xla/xla/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
2026-03-19 11:56:50.316213: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2026-03-19 11:56:51.821298: I external/local_xla/xla/tsl/cuda/cudart_stub.cc:31] Could not find cuda drivers on your machine, GPU will not be used.
>>> import maxtext
>>> import maxtext.configs
>>>

Unfortunately, the install_maxtext_tpu_github_deps does not quite work out-of-the-box in uv projects. It should be an easy fix on our end (replace "uv pip install" with "uv add"). However, it should not affect the basic import above.

[syrkis]

Very strange. Running your commands, including making a new uv project again, I am getting this:

desk/fun/clean_install on ☁️
❯ uv init --python 3.12
Initialized project `clean-install`

clean_install on  master [?] is 󰏗 v0.1.0 via  v3.14.3 on ☁️
❯ uv add "maxtext[tpu]>=0.2.0"
Using CPython 3.12.8 interpreter at: /usr/local/bin/python3.12
Creating virtual environment at: .venv
Resolved 263 packages in 3.35s
Prepared 14 packages in 13.70s
Installed 260 packages in 1.20s

clean_install on  master [?] is 󰏗 v0.1.0 via  v3.14.3 on ☁️   took 18s
❯ uv run python
Python 3.12.8 (v3.12.8:2dc476bcb91, Dec  3 2024, 14:43:19) [Clang 13.0.0 (clang-1300.0.29.30)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import maxtext
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'maxtext'
>>> import MaxText
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/nobr/desk/fun/clean_install/.venv/lib/python3.12/site-packages/MaxText/__init__.py", line 38, in <module>
    from maxtext.configs import pyconfig
ModuleNotFoundError: No module named 'maxtext'
>>>

[syrkis]

@igorts-git I reran this on my cluster on linux. There it works. I get the exact same as you. So the issue is on MacOS. Will find another mac to try to run it on to confirm

[syrkis]

I reran this on another mac, and I get the same issue:

>>> import maxtext
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'maxtext'
>>> import MaxText
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/syrkis/desk/fun/clean_install/.venv/lib/python3.12/site-packages/MaxText/__init__.py", line 38, in <module>
    from maxtext.configs import pyconfig
ModuleNotFoundError: No module named 'maxtext'
>>>

[igorts-git]

Unfortunately, I don't have a MacOs handy. My laptop is running Linux. I would love to better understand if the problem is due to your usage of a uv managed project (with pyproject.toml).
Could you make a new clean directory and try to use the "classic" venv commands as below instead of "uv init" and "uv add"?

uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate
uv pip install "maxtext[tpu]>=0.2.0" --resolution=lowest
install_maxtext_tpu_github_deps

Thanks!

[syrkis]

Just checked. The issue remains when installing with the commands above in a new clean directory (no uv init or uv add). I ran all four commands, but I've truncated the cli output to the last command, and a test to see if the issue remains. It does.

desk/fun/clean_install via  v3.12.8 (maxtext_venv) on ☁️   took 2m14s
❯ install_maxtext_tpu_github_deps
Installing extra dependencies from '/Users/nobr/desk/fun/clean_install/maxtext_venv/lib/python3.12/site-packages/install_maxtext_extra_deps/extra_deps_from_github.txt' using uv...
Running command: /Users/nobr/desk/fun/clean_install/maxtext_venv/bin/python3 -m uv pip install -r /Users/nobr/desk/fun/clean_install/maxtext_venv/lib/python3.12/site-packages/install_maxtext_extra_deps/extra_deps_from_github.txt --no-deps
Extra dependencies installed successfully!
--- Output from uv ---

--- Errors/Warnings from uv (if any) ---
Using Python 3.12.8 environment at: maxtext_venv
Checked 2 packages in 27ms


desk/fun/clean_install via  v3.12.8 (maxtext_venv) on ☁️
❯ uv run python
Python 3.12.8 (v3.12.8:2dc476bcb91, Dec  3 2024, 14:43:19) [Clang 13.0.0 (clang-1300.0.29.30)] on darwin
Type "help", "copyright", "credits" or "license" for more information.
>>> import MaxText
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/nobr/desk/fun/clean_install/maxtext_venv/lib/python3.12/site-packages/MaxText/__init__.py", line 38, in <module>
    from maxtext.configs import pyconfig
ModuleNotFoundError: No module named 'maxtext'
>>>

[syrkis]

@igorts-git I'll make a fork and try to fix it tomorrow :)