Move layers and models into new folders

Author: dipannita08

src/MaxText/__init__.py

@@ -36,7 +36,7 @@
 from jax.sharding import Mesh
 
 from MaxText import pyconfig
-from MaxText.layers import models
+from maxtext.models import models
 from maxtext.trainers.post_train.dpo import dpo_utils
 from maxtext.utils import maxtext_utils
 from maxtext.utils import model_creation_utils

src/MaxText/generate_param_only_checkpoint.py

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

src/MaxText/integration/tunix/tunix_adapter.py

@@ -25,7 +25,7 @@
 
 from jax import Array
 from flax import nnx
-from MaxText.layers.models import Transformer
+from maxtext.models.models import Transformer
 from MaxText.integration.tunix.utils import VllmWeightMapping
 from maxtext.checkpoint_conversion.utils.hf_model_configs import HF_MODEL_CONFIGS  # pylint: disable=ungrouped-imports
 

src/MaxText/layerwise_quantization.py

@@ -45,7 +45,8 @@
 
 from MaxText import common_types
 from MaxText import pyconfig
-from MaxText.layers import models, quantizations, deepseek
+from maxtext.models import models, deepseek
+from maxtext.layers import quantizations
 from maxtext.common import checkpointing
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils

src/MaxText/maxengine.py

@@ -39,7 +39,8 @@
 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.layers import models, quantizations
+from maxtext.models import models
+from maxtext.layers import quantizations
 from maxtext.inference import inference_utils
 from maxtext.inference.page_manager import PageManager, PageState
 from maxtext.multimodal import processor as mm_processor

src/MaxText/train.py

@@ -40,15 +40,14 @@
 
 from MaxText import pyconfig
 from MaxText import sharding
-from MaxText.layers.multi_token_prediction import calculate_mtp_acceptance_rate, calculate_mtp_loss
 from MaxText.common_types import ShardMode
 from MaxText.globals import EPS
 # Placeholder: internal
 
 from MaxText.gradient_accumulation import gradient_accumulation_loss_and_grad
 from MaxText.vocabulary_tiling import vocab_tiling_linen_loss
 # pylint: disable=too-many-positional-arguments
-
+from maxtext.layers.multi_token_prediction import calculate_mtp_acceptance_rate, calculate_mtp_loss
 from maxtext.common import checkpointing, profiler
 from maxtext.common.goodput import (
     GoodputEvent,

src/MaxText/train_compile.py

@@ -41,8 +41,8 @@
 from MaxText import pyconfig
 from MaxText import sharding
 from MaxText.common_types import MODEL_MODE_TRAIN, ShardMode
-from MaxText.layers import models
-from MaxText.layers import quantizations
+from maxtext.models import models
+from maxtext.layers import quantizations
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils

src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py

@@ -53,8 +53,8 @@
 from MaxText import pyconfig
 from MaxText.common_types import MODEL_MODE_TRAIN
 from MaxText.globals import MAXTEXT_PKG_DIR
-from MaxText.layers import quantizations
-from MaxText.layers.models import transformer_as_linen
+from maxtext.layers import quantizations
+from maxtext.models.models import transformer_as_linen
 from maxtext.common import checkpointing
 from maxtext.utils import max_logging
 from maxtext.utils import maxtext_utils

src/maxtext/checkpoint_conversion/to_maxtext.py

@@ -77,11 +77,12 @@
 from orbax.checkpoint import type_handlers
 from MaxText import pyconfig
 from MaxText.common_types import MODEL_MODE_TRAIN
-from MaxText.layers import models, quantizations
 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 apply_hook_fns, HF_IDS, print_ram_usage, get_hf_model, MemoryMonitorTqdm, print_peak_memory, validate_and_filter_param_map_keys
 from maxtext.inference.inference_utils import str2bool
+from maxtext.models import models
+from maxtext.layers import quantizations
 from maxtext.utils import max_logging
 from maxtext.utils import max_utils
 from maxtext.utils import maxtext_utils

src/maxtext/inference/kvcache.py

@@ -26,8 +26,8 @@
 from aqt.jax.v2.aqt_tensor import QTensor as KVTensor
 from aqt.jax.v2.flax import aqt_flax
 
-from MaxText.layers import nnx_wrappers
-from MaxText.layers.initializers import variable_to_logically_partitioned
+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