Merge pull request #3026 from AI-Hypercomputer:utils_refactor

PiperOrigin-RevId: 862910344

Author: Google-ML-Automation

benchmarks/api_server/maxtext_generator.py

@@ -34,7 +34,8 @@
 
 from dataclasses import dataclass, field
 
-from MaxText import max_utils, maxengine, pyconfig, multimodal_utils, max_logging
+from MaxText import maxengine, pyconfig, multimodal_utils
+from maxtext.utils import max_logging, max_utils
 
 # Set TF log level to avoid verbose startup messages.
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

benchmarks/mmlu/mmlu_eval.py

@@ -57,9 +57,9 @@
 from tqdm import tqdm
 
 from MaxText import pyconfig
-from MaxText import max_logging
-from MaxText import max_utils
 from MaxText import maxengine
+from maxtext.utils import max_logging
+from maxtext.utils import max_utils
 
 ASCII_UPPERCASE_A = ord("A")  # ASCII value for uppercase 'A'
 

src/MaxText/__init__.py

@@ -29,12 +29,12 @@
 
 from jax.sharding import Mesh
 
-from MaxText import maxtext_utils
-from MaxText import model_creation_utils
 from MaxText import pyconfig
 from MaxText.layers import models
 from MaxText import dpo_utils
-from MaxText.model_creation_utils import from_config
+from maxtext.utils import maxtext_utils
+from maxtext.utils import model_creation_utils
+from maxtext.utils.model_creation_utils import from_config
 
 Transformer = models.Transformer
 transformer_as_linen = models.transformer_as_linen

src/MaxText/benchmark_chunked_prefill.py

@@ -47,9 +47,9 @@
 
 from absl import app
 
-from MaxText import max_utils
 from MaxText import maxengine
 from MaxText import pyconfig
+from maxtext.utils import max_utils
 
 _WARMUP_ITERS = 2
 _BENCHMARK_ITERS = 5

src/MaxText/configs/types.py

@@ -33,10 +33,11 @@
 from pydantic.main import BaseModel
 from pydantic.types import PositiveInt, NonNegativeFloat, NonNegativeInt
 
-from MaxText import accelerator_to_spec_map, max_utils
+from MaxText import accelerator_to_spec_map
 from MaxText.common_types import AttentionType, DecoderBlockType, ShardMode
 from MaxText.globals import MAXTEXT_ASSETS_ROOT
-from MaxText.utils import gcs_utils
+from maxtext.utils import gcs_utils
+from maxtext.utils import max_utils
 
 logger = logging.getLogger(__name__)
 

src/MaxText/decode.py

@@ -23,12 +23,12 @@
 
 from jetstream.engine import engine_api
 
-from MaxText import max_utils
 from MaxText import maxengine
 from MaxText import pyconfig
 from MaxText import multimodal_utils
 from MaxText.multimodal import preprocessor
 from maxtext.common import profiler
+from maxtext.utils import max_utils
 # Placeholder: internal
 
 # Number of text sequences to process in a single batch.

src/MaxText/distillation/train_distill.py

@@ -17,7 +17,7 @@
 import sys
 import importlib
 
-from MaxText import max_logging
+from maxtext.utils import max_logging
 
 OLD_MODULE_PATH = "MaxText.distillation.train_distill"
 NEW_MODULE_PATH = "maxtext.trainers.post_train.distillation.train_distill"

src/MaxText/dpo_utils.py

@@ -19,7 +19,7 @@
 import jax
 import jax.numpy as jnp
 
-from MaxText import maxtext_utils
+from maxtext.utils import maxtext_utils
 
 
 def _split_dpo_state(state):

src/MaxText/examples/demo_decoding.ipynb

@@ -135,19 +135,19 @@
         "\n",
         "import MaxText as mt\n",
         "from MaxText import common_types\n",
-        "from MaxText import maxtext_utils\n",
-        "from MaxText import max_logging\n",
         "from MaxText import pyconfig\n",
         "from MaxText.input_pipeline import _input_pipeline_utils\n",
         "from MaxText.utils.ckpt_conversion import to_maxtext\n",
-        "from maxtext.src.maxtext.inference import inference_utils\n",
+        "from maxtext.inference import inference_utils\n",
+        "from maxtext.utils import maxtext_utils\n",
+        "from maxtext.utils import max_logging\n",
         "\n",
         "from google.colab import userdata\n",
         "from huggingface_hub import login\n",
         "\n",
         "MAXTEXT_PKG_DIR = os.path.dirname(mt.__file__)\n",
-    "MAXTEXT_REPO_ROOT = os.path.dirname(os.path.dirname(MAXTEXT_PKG_DIR))\n",
-    "MAXTEXT_ASSETS_ROOT = os.path.join(MAXTEXT_REPO_ROOT, \"src\", \"maxtext\", \"assets\")\n",
+        "MAXTEXT_REPO_ROOT = os.path.dirname(os.path.dirname(MAXTEXT_PKG_DIR))\n",
+        "MAXTEXT_ASSETS_ROOT = os.path.join(MAXTEXT_REPO_ROOT, \"src\", \"maxtext\", \"assets\")\n",
         "\n",
         "nest_asyncio.apply()"
       ]