Internal only

PiperOrigin-RevId: 829542624

Author: abhinavclemson

src/MaxText/integration/tunix/tunix_adapter.py

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

src/MaxText/integration/tunix/weight_mapping/__init__.py

@@ -19,8 +19,8 @@
 model name. This allows for easy extension to support new models.
 """
 
-from maxtext.src.maxtext.integration.tunix.weight_mapping.llama3 import LLAMA3_VLLM_MAPPING
-from maxtext.src.maxtext.integration.tunix.weight_mapping.qwen3 import QWEN3_VLLM_MAPPING
+from MaxText.integration.tunix.weight_mapping.llama3 import LLAMA3_VLLM_MAPPING
+from MaxText.integration.tunix.weight_mapping.qwen3 import QWEN3_VLLM_MAPPING
 
 
 class StandaloneVllmWeightMapping:

src/MaxText/rl/train_rl.py

@@ -75,7 +75,10 @@
 os.environ["SKIP_JAX_PRECOMPILE"] = "1"
 
 
-from MaxText import max_logging, max_utils, maxtext_utils, pyconfig
+from MaxText import max_logging
+from MaxText import max_utils
+from MaxText import maxtext_utils
+from MaxText import pyconfig
 from MaxText import model_creation_utils
 from MaxText.integration.tunix.tunix_adapter import TunixMaxTextAdapter
 from MaxText.rl.evaluate_rl import evaluate