pre-nnx-v0.1.0

Use this release for the latest MaxText version that fully depends on Flax Linen (no NNX).

tpu-recipes-v0.1.3

Use this release for tpu-recipes that require version tpu-recipes-v0.1.3

tpu-recipes-v0.1.2

Use this release for tpu-recipes that require version tpu-recipes-v0.1.2

tpu-recipes-v0.1.1

Use this release for tpu-recipes that require version tpu-recipes-v0.1.1

pre-module-v0.1.0

Release prior to module refactor for older train API.

With this release or prior: python3 MaxText/train.py MaxText/configs/base.yml run_name=...

Soon after this release: python3 -m MaxText.train MaxText/configs/base.yml run_name=...

tpu-recipes-v0.1.0

Use this release for tpu-recipes that require version tpu-recipes-v0.1.0

MoE v1.0.0

MoE v1.0.0 supports:

• Megablox with Fully Sharded Data Parallelism (FSDP) and Token Parallelism (TP)
• Dropping strategies with FSDP, TP, and Expert Parallelism (EP)