Merge pull request #3278 from AI-Hypercomputer:bvandermoon-maxtext-restructure

PiperOrigin-RevId: 876541062

Author: Google-ML-Automation

src/MaxText/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023–2025 Google LLC
+# Copyright 2023–2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

src/maxtext/experimental/rl/grpo_trainer.py

@@ -66,7 +66,7 @@
 
 from ml_goodput_measurement.src.goodput import GoodputRecorder
 
-import MaxText as mt
+import maxtext as mt
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import EPS
 from maxtext.trainers.pre_train.train import get_first_step

tests/integration/grpo_trainer_correctness_test.py

@@ -34,7 +34,7 @@
 import jax.numpy as jnp
 from jax.sharding import Mesh
 import jsonlines
-import MaxText as mt
+import maxtext as mt
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT, MAXTEXT_PKG_DIR, MAXTEXT_TEST_ASSETS_ROOT
 from maxtext.common.common_types import MODEL_MODE_TRAIN

tests/unit/engram_vs_reference_test.py

@@ -46,7 +46,7 @@
 from jax.sharding import Mesh
 
 from maxtext.configs import pyconfig
-from MaxText import maxtext_utils
+from maxtext import maxtext_utils
 
 from maxtext.layers.engram import CompressedTokenizer as CompressedTokenizerJAX
 from maxtext.layers.engram import NgramHashMapping as NgramHashMappingJAX

tests/unit/moba_vs_reference_test.py

@@ -31,7 +31,7 @@
 import torch
 from jax.sharding import Mesh
 
-from MaxText import maxtext_utils, pyconfig
+from maxtext import maxtext_utils, pyconfig
 from maxtext.layers.attention_op import AttentionOp
 from tests.utils.test_helpers import get_test_config_path
 

tests/unit/sharding_compare_test.py

@@ -19,7 +19,7 @@
 import os
 import jax
 import jax.numpy as jnp
-from MaxText import maxtext_utils
+from maxtext import maxtext_utils
 from maxtext.configs import pyconfig
 # import optax
 

tests/utils/sharding_dump.py

@@ -26,7 +26,7 @@
 import jax
 from jax.sharding import NamedSharding, PartitionSpec
 from jax.tree_util import tree_flatten_with_path
-from MaxText import maxtext_utils
+from maxtext import maxtext_utils
 from maxtext.configs import pyconfig
 
 from maxtext.utils.globals import MAXTEXT_REPO_ROOT