Fix Broken MaxText CI

Author: dipannita08

src/maxtext/trainers/pre_train/train.py

@@ -42,8 +42,6 @@
 from MaxText import sharding
 from MaxText.common_types import ShardMode
 from MaxText.globals import EPS
-# pylint: disable-next=unused-import
-from maxtext import maxtext_google
 
 from MaxText.gradient_accumulation import gradient_accumulation_loss_and_grad
 from MaxText.vocabulary_tiling import vocab_tiling_linen_loss

tests/unit/llama4_layers_test.py

@@ -27,7 +27,8 @@
 
 from MaxText.common_types import MODEL_MODE_TRAIN, AttentionType
 from MaxText import pyconfig
-from maxtext.layers import attentions, embeddings, llama4
+from maxtext.layers import attentions, embeddings
+from maxtext.models import llama4
 from maxtext.utils import maxtext_utils
 import numpy as np
 from tests.utils.test_helpers import get_test_config_path

tests/unit/pipeline_parallelism_test.py

@@ -26,13 +26,16 @@
 import jax.numpy as jnp
 from jax.sharding import Mesh
 from MaxText import pyconfig
-from maxtext.common.gcloud_stub import is_decoupled
 from MaxText.common_types import MODEL_MODE_TRAIN
 from MaxText.globals import MAXTEXT_ASSETS_ROOT
+from maxtext.common.gcloud_stub import is_decoupled
 from maxtext.layers import nnx_wrappers
 from maxtext.layers import pipeline
-from maxtext.layers import simple_layer
+from maxtext.models import deepseek
+from maxtext.models import simple_layer
+from maxtext.utils import maxtext_utils
 from maxtext.trainers.pre_train.train import main as train_main
+from tests.utils.test_helpers import get_test_config_path, get_test_dataset_path, get_test_base_output_directory
 import pytest