Migrate maxtext common_types

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/maxengine.py

@@ -37,7 +37,6 @@
 import flax
 
 from MaxText import pyconfig
-from MaxText.common_types import MODEL_MODE_PREFILL, DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_AUTOREGRESSIVE
 from MaxText.globals import MAXTEXT_PKG_DIR
 from maxtext.models import models
 from maxtext.layers import quantizations
@@ -48,6 +47,7 @@
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils
 from maxtext.common.gcloud_stub import jetstream, is_decoupled
+from maxtext.common.common_types import MODEL_MODE_PREFILL, DECODING_ACTIVE_SEQUENCE_INDICATOR, MODEL_MODE_AUTOREGRESSIVE
 
 config_lib, engine_api, token_utils, tokenizer_api, _token_params_ns = jetstream()
 TokenizerParameters = getattr(_token_params_ns, "TokenizerParameters", object)  # type: ignore[assignment]

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

@@ -74,7 +74,7 @@
 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
-from MaxText.common_types import MODEL_MODE_TRAIN
+from maxtext.common.common_types import MODEL_MODE_TRAIN
 from maxtext.inference.inference_utils import str2bool
 from maxtext.layers import quantizations
 from maxtext.models import models

src/MaxText/common_types.py → 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.

src/maxtext/configs/types.py

@@ -27,18 +27,17 @@
 from typing import Any, Literal, NewType, Optional
 
 import jax
-from MaxText import accelerator_to_spec_map
-from MaxText.common_types import AttentionType, DecoderBlockType, ShardMode
-from MaxText.globals import MAXTEXT_ASSETS_ROOT
+from maxtext.common.common_types import AttentionType, DecoderBlockType, ShardMode
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_utils
+from MaxText import accelerator_to_spec_map
+from MaxText.globals import MAXTEXT_ASSETS_ROOT
 from pydantic.config import ConfigDict
 from pydantic.fields import Field
 from pydantic.functional_validators import field_validator, model_validator
 from pydantic.main import BaseModel
 from pydantic.types import NonNegativeFloat, NonNegativeInt, PositiveInt
 
-
 class XProfTPUPowerTraceMode(enum.IntEnum):  # pylint: disable=invalid-name
   """Enum for XProfTPUPowerTraceMode."""
 

src/maxtext/examples/demo_decoding.ipynb

@@ -134,7 +134,7 @@
         "import numpy as np\n",
         "\n",
         "import MaxText as mt\n",
-        "from MaxText import common_types\n",
+        "from maxtext.common import common_types\n",
         "from MaxText import pyconfig\n",
         "from MaxText.input_pipeline import _input_pipeline_utils\n",
         "from maxtext.checkpoint_conversion import to_maxtext\n",

src/maxtext/experimental/rl/grpo_utils.py

@@ -21,7 +21,7 @@
 import jaxtyping
 from typing import Any, Callable
 
-from MaxText.common_types import DecoderBlockType
+from maxtext.common.common_types import DecoderBlockType
 from maxtext.inference.offline_engine import InputData
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils

src/maxtext/inference/kvcache.py

@@ -29,8 +29,8 @@
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.initializers import variable_to_logically_partitioned
 
-from MaxText.common_types import Array, AxisNames, AxisIdxes, Config, CACHE_BATCH_PREFILL, DType, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, MODEL_MODE_AUTOREGRESSIVE, CACHE_HEADS_NONE, DECODING_ACTIVE_SEQUENCE_INDICATOR
-from MaxText.common_types import CACHE_BATCH, CACHE_SEQUENCE, CACHE_HEADS, CACHE_KV, CACHE_SCALE_BATCH, CACHE_SCALE_SEQUENCE, CACHE_SCALE_HEADS, CACHE_SCALE_KV
+from maxtext.common.common_types import Array, AxisNames, AxisIdxes, Config, CACHE_BATCH_PREFILL, DType, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, MODEL_MODE_AUTOREGRESSIVE, CACHE_HEADS_NONE, DECODING_ACTIVE_SEQUENCE_INDICATOR
+from maxtext.common.common_types import CACHE_BATCH, CACHE_SEQUENCE, CACHE_HEADS, CACHE_KV, CACHE_SCALE_BATCH, CACHE_SCALE_SEQUENCE, CACHE_SCALE_HEADS, CACHE_SCALE_KV
 
 
 MAX_INT8 = 127.5

src/maxtext/inference/page_manager.py

@@ -31,7 +31,7 @@
 
 from jaxtyping import Array, Integer, Bool
 
-from MaxText.common_types import Config
+from maxtext.common.common_types import Config
 
 # Aliases using <Dims><Type><Rank>d convention
 # We use string names for dimensions as they are symbolic within the type hints.

src/maxtext/inference/paged_attention.py

@@ -26,7 +26,7 @@
 import jax.numpy as jnp
 from jax.sharding import Mesh
 from jax.sharding import PartitionSpec as P
-from MaxText.common_types import Array, AxisNames, BATCH, DType, D_KV, HEAD, LENGTH, MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL
+from maxtext.common.common_types import Array, AxisNames, BATCH, DType, D_KV, HEAD, LENGTH, MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL
 from maxtext.inference import page_manager
 from maxtext.inference import paged_attention_kernel_v2
 from maxtext.layers.initializers import variable_to_logically_partitioned

src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py

@@ -22,8 +22,8 @@
 from jax import numpy as jnp
 from jax.sharding import Mesh
 from MaxText import pyconfig
-from MaxText.common_types import MODEL_MODE_AUTOREGRESSIVE
 from MaxText.globals import MAXTEXT_CONFIGS_DIR
+from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE
 from maxtext.utils import max_logging
 from maxtext.utils import model_creation_utils
 

src/maxtext/kernels/attention/ragged_attention.py

@@ -24,7 +24,7 @@
 from jax.experimental.pallas import tpu as pltpu
 import jax.numpy as jnp
 
-from MaxText.common_types import DEFAULT_MASK_VALUE
+from maxtext.common.common_types import DEFAULT_MASK_VALUE
 
 
 def get_mha_cost_estimate(shape_dtype):

src/maxtext/layers/attention_mla.py

@@ -32,7 +32,7 @@
 
 from flax import nnx
 
-from MaxText.common_types import (
+from maxtext.common.common_types import (
     Array,
     AxisIdxes,
     AxisNames,

src/maxtext/layers/attention_op.py

@@ -33,7 +33,7 @@
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_mask
 import jax.numpy as jnp
 from jax.sharding import Mesh, NamedSharding
-from MaxText.common_types import (
+from maxtext.common.common_types import (
     Array,
     AttentionType,
     AxisIdxes,

src/maxtext/layers/attentions.py

@@ -25,7 +25,7 @@
 
 from flax import nnx
 
-from MaxText.common_types import (
+from maxtext.common.common_types import (
     DecoderBlockType,
     BATCH,
     BATCH_NO_EXP,

src/maxtext/layers/decoders.py

@@ -27,8 +27,8 @@
 import jax.numpy as jnp
 from jax.sharding import Mesh
 from MaxText import sharding
-from MaxText.common_types import Config, DecoderBlockType, EP_AS_CONTEXT, ShardMode
-from MaxText.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN
+from maxtext.common.common_types import Config, DecoderBlockType, EP_AS_CONTEXT, ShardMode
+from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN
 from maxtext.inference import page_manager
 from maxtext.layers import linears
 from maxtext.layers import mhc

src/maxtext/layers/embeddings.py

@@ -25,7 +25,7 @@
 from flax import nnx
 
 from MaxText.sharding import logical_to_mesh_axes, create_sharding
-from MaxText.common_types import ShardMode, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, Array, Config, DType
+from maxtext.common.common_types import ShardMode, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN, Array, Config, DType
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.initializers import Initializer, default_embed_init, variable_to_logically_partitioned
 from maxtext.utils import max_logging

src/maxtext/layers/encoders.py

@@ -18,7 +18,7 @@
 from flax import nnx
 from jax.sharding import Mesh
 
-from MaxText.common_types import Config
+from maxtext.common.common_types import Config
 from maxtext.layers import nnx_wrappers
 from maxtext.layers import initializers
 

src/maxtext/layers/engram.py

@@ -24,7 +24,7 @@
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh
-from MaxText.common_types import Array, Config, MODEL_MODE_TRAIN
+from maxtext.common.common_types import Array, Config, MODEL_MODE_TRAIN
 from maxtext.input_pipeline.tokenizer import HFTokenizer
 from maxtext.layers.embeddings import Embed
 from maxtext.layers.initializers import NdInitializer, nd_dense_init

src/maxtext/layers/initializers.py

@@ -22,7 +22,7 @@
 from flax import nnx
 from aqt.jax.v2 import aqt_tensor
 
-from MaxText.common_types import Array, DType, Shape, PRNGKey
+from maxtext.common.common_types import Array, DType, Shape, PRNGKey
 
 Initializer = Callable[[PRNGKey, Shape, DType], Array]
 InitializerAxis = int | tuple[int, ...]

src/maxtext/layers/linears.py

@@ -30,8 +30,8 @@
 import flax.linen as nn
 
 from MaxText.sharding import maybe_shard_with_logical
-from MaxText.common_types import DecoderBlockType, ShardMode, DType, Array, Config
-from MaxText.common_types import MODEL_MODE_TRAIN, MODEL_MODE_PREFILL, EP_AS_CONTEXT
+from maxtext.common.common_types import DecoderBlockType, ShardMode, DType, Array, Config
+from maxtext.common.common_types import MODEL_MODE_TRAIN, MODEL_MODE_PREFILL, EP_AS_CONTEXT
 from maxtext.layers import nnx_wrappers, quantizations
 from maxtext.layers import normalizations
 from maxtext.layers.initializers import NdInitializer, nd_dense_init, default_bias_init, variable_to_logically_partitioned

src/maxtext/layers/mhc.py

@@ -19,8 +19,8 @@
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh
-from MaxText.common_types import Array, Config
-from MaxText.common_types import HyperConnectionType
+from maxtext.common.common_types import Array, Config
+from maxtext.common.common_types import HyperConnectionType
 from maxtext.layers.initializers import default_bias_init, default_scalar_init, nd_dense_init
 from maxtext.layers.normalizations import RMSNorm
 

src/maxtext/layers/moe.py

@@ -29,8 +29,8 @@
 from jax.sharding import NamedSharding, Mesh
 from jax.sharding import PartitionSpec as P
 import jax.numpy as jnp
-from MaxText import common_types as ctypes
-from MaxText.common_types import ShardMode
+from maxtext.common import common_types as ctypes
+from maxtext.common.common_types import ShardMode
 from MaxText.sharding import maybe_shard_with_logical, create_sharding
 from MaxText.sharding import logical_to_mesh_axes
 from maxtext.layers import attentions, linears, nnx_wrappers, quantizations

src/maxtext/layers/multi_token_prediction.py

@@ -22,7 +22,7 @@
 import jax.numpy as jnp
 from jax.sharding import Mesh
 from MaxText import sharding
-from MaxText.common_types import Config, MODEL_MODE_TRAIN
+from maxtext.common.common_types import Config, MODEL_MODE_TRAIN
 from MaxText.globals import EPS
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.decoders import DecoderLayer

src/maxtext/layers/normalizations.py

@@ -23,7 +23,7 @@
 from jax import lax
 import jax.numpy as jnp
 from jax.sharding import NamedSharding
-from MaxText.common_types import Array, DType, ShardMode
+from maxtext.common.common_types import Array, DType, ShardMode
 from maxtext.layers import nnx_wrappers
 from maxtext.layers.initializers import Initializer, variable_to_logically_partitioned
 from maxtext.utils import max_logging

src/maxtext/layers/pipeline.py

@@ -28,7 +28,7 @@
 from flax import linen as nn
 from flax.linen.spmd import LogicallyPartitioned
 
-from MaxText.common_types import Config, MODEL_MODE_TRAIN, EP_AS_CONTEXT, ShardMode
+from maxtext.common.common_types import Config, MODEL_MODE_TRAIN, EP_AS_CONTEXT, ShardMode
 from MaxText.sharding import (
     maybe_shard_with_logical,
     maybe_shard_with_name,

src/maxtext/layers/quantizations.py

@@ -36,7 +36,7 @@
 from flax.linen import initializers as flax_initializers
 import flax.linen as nn
 
-from MaxText.common_types import DType, Config
+from maxtext.common.common_types import DType, Config
 from maxtext.inference.kvcache import KVQuant
 
 # Params used to define mixed precision quantization configs

src/maxtext/models/deepseek.py

@@ -22,8 +22,8 @@
 from jax.ad_checkpoint import checkpoint_name
 import jax.numpy as jnp
 from jax.sharding import Mesh
-from MaxText.common_types import Config
-from MaxText.common_types import HyperConnectionType, MODEL_MODE_PREFILL
+from maxtext.common.common_types import Config
+from maxtext.common.common_types import HyperConnectionType, MODEL_MODE_PREFILL
 from maxtext.inference import page_manager
 from maxtext.layers import attention_mla
 from maxtext.layers import initializers

src/maxtext/models/gemma.py

@@ -22,7 +22,7 @@
 from jax.sharding import Mesh
 import jax.numpy as jnp
 
-from MaxText.common_types import Config
+from maxtext.common.common_types import Config
 from maxtext.layers import initializers
 from maxtext.layers import nnx_wrappers
 from maxtext.layers import quantizations

src/maxtext/models/gemma2.py

@@ -22,7 +22,7 @@
 from jax.sharding import Mesh
 import jax.numpy as jnp
 
-from MaxText.common_types import MODEL_MODE_PREFILL, Config
+from maxtext.common.common_types import MODEL_MODE_PREFILL, Config
 from maxtext.layers import attentions
 from maxtext.layers import initializers
 from maxtext.layers import nnx_wrappers

src/maxtext/models/gemma3.py

@@ -22,7 +22,7 @@
 from flax import linen as nn
 from flax import nnx
 
-from MaxText.common_types import Config, AttentionType, MODEL_MODE_PREFILL
+from maxtext.common.common_types import Config, AttentionType, MODEL_MODE_PREFILL
 from maxtext.layers import quantizations
 from maxtext.layers import nnx_wrappers
 from maxtext.layers import initializers

src/maxtext/models/gpt3.py

@@ -27,7 +27,7 @@
 from flax import linen as nn
 from flax import nnx
 
-from MaxText.common_types import Config, DType, AxisNames, BATCH, LENGTH, EMBED, HEAD, D_KV, Array, MODEL_MODE_TRAIN
+from maxtext.common.common_types import Config, DType, AxisNames, BATCH, LENGTH, EMBED, HEAD, D_KV, Array, MODEL_MODE_TRAIN
 from maxtext.layers import initializers, nnx_wrappers
 from maxtext.layers.linears import DenseGeneral, MlpBlock, canonicalize_tuple, normalize_axes
 from maxtext.layers import quantizations