Update import paths for maxtext_utils and pyconfig in tests

PiperOrigin-RevId: 877446791

Author: SurbhiJainUSC

src/install_maxtext_extra_deps/install_post_train_extra_deps.py

@@ -71,7 +71,7 @@ def main():
         "uv",
         "pip",
         "install",
-        "src/maxtext/integration/vllm",
+        "src/MaxText/integration/vllm",
         "--no-deps",
   ]
 

tests/unit/engram_vs_reference_test.py

@@ -46,13 +46,12 @@
 from jax.sharding import Mesh
 
 from maxtext.configs import pyconfig
-from maxtext import maxtext_utils
-
 from maxtext.layers.engram import CompressedTokenizer as CompressedTokenizerJAX
 from maxtext.layers.engram import NgramHashMapping as NgramHashMappingJAX
 from maxtext.layers.engram import MultiHeadEmbedding as MultiHeadEmbeddingJAX
 from maxtext.layers.engram import ShortConv as ShortConvJAX
 from maxtext.layers.engram import Engram as EngramJAX
+from maxtext.utils import maxtext_utils
 from maxtext.utils.globals import MAXTEXT_PKG_DIR
 
 

tests/unit/moba_vs_reference_test.py

@@ -31,8 +31,9 @@
 import torch
 from jax.sharding import Mesh
 
-from maxtext import maxtext_utils, pyconfig
+from maxtext.configs import pyconfig
 from maxtext.layers.attention_op import AttentionOp
+from maxtext.utils import maxtext_utils
 from tests.utils.test_helpers import get_test_config_path
 
 # pylint: disable=missing-function-docstring,protected-access

tests/unit/sharding_compare_test.py

@@ -19,8 +19,8 @@
 import os
 import jax
 import jax.numpy as jnp
-from maxtext import maxtext_utils
 from maxtext.configs import pyconfig
+from maxtext.utils import maxtext_utils
 # import optax
 
 from maxtext.utils.globals import MAXTEXT_PKG_DIR

tests/utils/sharding_dump.py

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