Merge pull request #3205 from AI-Hypercomputer:move-maxtext-common

PiperOrigin-RevId: 874806284

Author: Google-ML-Automation

src/MaxText/generate_param_only_checkpoint.py

@@ -32,7 +32,7 @@
 from MaxText import optimizers
 from MaxText import pyconfig
 from maxtext.common import checkpointing
-from MaxText.common_types import DecoderBlockType, MODEL_MODE_TRAIN
+from maxtext.common.common_types import DecoderBlockType, MODEL_MODE_TRAIN
 from maxtext.layers import quantizations
 from maxtext.models import models
 from maxtext.utils import gcs_utils

src/MaxText/gradient_accumulation.py

@@ -18,7 +18,7 @@
 import jax.numpy as jnp
 from jax.sharding import NamedSharding
 
-from MaxText.common_types import ShardMode
+from maxtext.common.common_types import ShardMode
 from MaxText.sharding import maybe_shard_with_name
 
 

src/MaxText/layerwise_quantization.py

@@ -39,8 +39,7 @@
 from flax.linen import partitioning as nn_partitioning
 import jax
 import jax.numpy as jnp
-from MaxText import common_types
-from MaxText import pyconfig
+from maxtext.common import common_types
 from maxtext.common import checkpointing
 from maxtext.layers import quantizations
 from maxtext.models import deepseek, models
@@ -49,6 +48,7 @@
 from maxtext.utils import maxtext_utils
 import orbax.checkpoint as ocp
 from tqdm import tqdm
+from MaxText import pyconfig
 
 IGNORE = ocp.PLACEHOLDER
 PRNGKeyType = Any

src/MaxText/pyconfig.py

@@ -29,8 +29,8 @@
 import omegaconf
 
 from MaxText import pyconfig_deprecated
-from MaxText.common_types import DecoderBlockType, ShardMode
 from MaxText.globals import MAXTEXT_CONFIGS_DIR
+from maxtext.common.common_types import DecoderBlockType, ShardMode
 from maxtext.configs import types
 from maxtext.configs.types import MaxTextConfig
 from maxtext.inference.inference_utils import str2bool

src/MaxText/pyconfig_deprecated.py

@@ -30,7 +30,7 @@
 
 from MaxText import accelerator_to_spec_map
 from MaxText.globals import MAXTEXT_ASSETS_ROOT, MAXTEXT_REPO_ROOT, MAXTEXT_PKG_DIR
-from MaxText.common_types import AttentionType, DecoderBlockType, ShardMode
+from maxtext.common.common_types import AttentionType, DecoderBlockType, ShardMode
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils

src/MaxText/sharding.py

@@ -25,7 +25,7 @@
 
 import optax
 
-from MaxText.common_types import ShardMode
+from maxtext.common.common_types import ShardMode
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 

src/MaxText/vocabulary_tiling.py

@@ -25,7 +25,7 @@
     all_gather_over_fsdp,
     create_sharding,
 )
-from MaxText.common_types import ShardMode
+from maxtext.common.common_types import ShardMode
 from maxtext.utils import max_utils
 
 

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -44,9 +44,9 @@
 from jax.sharding import Mesh
 from MaxText import optimizers
 from MaxText import pyconfig
-from maxtext.common import checkpointing
-from MaxText.common_types import MODEL_MODE_TRAIN
 from MaxText.globals import MAXTEXT_PKG_DIR
+from maxtext.common import checkpointing
+from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.layers import quantizations
 from maxtext.models.models import transformer_as_linen
 from maxtext.utils import max_logging

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -71,7 +71,7 @@
 from huggingface_hub import hf_hub_download, list_repo_files
 import jax
 from MaxText import pyconfig
-from MaxText.common_types import MODEL_MODE_TRAIN
+from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.checkpoint_conversion.standalone_scripts.llama_or_mistral_ckpt import save_weights_to_checkpoint
 from maxtext.checkpoint_conversion.utils.param_mapping import HOOK_FNS, PARAM_MAPPING
 from maxtext.checkpoint_conversion.utils.utils import HF_IDS, MemoryMonitorTqdm, apply_hook_fns, get_hf_model, print_peak_memory, print_ram_usage, validate_and_filter_param_map_keys

src/MaxText/common_types.py src/maxtext/common/common_types.pysrc/MaxText/common_types.py renamed to src/maxtext/common/common_types.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.