Fix rl and integration imports

PiperOrigin-RevId: 829620709

Author: gobbleturk

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

@@ -19,7 +19,7 @@
 from tqdm.auto import tqdm
 from tunix.rl.rollout.base_rollout import RolloutConfig
 
-from MaxText.rl import utils_rl
+from  maxtext.src.maxtext.rl import utils_rl
 from MaxText import max_logging
 
 # ## Evaluate

src/MaxText/rl/train_rl.py

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