Import fix for E2E tests

Author: SurbhiJainUSC

end_to_end/tpu/gemma2/2b/test_gemma2_to_mt.sh

@@ -64,7 +64,7 @@ python3 -m MaxText.decode "${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/Max
 # Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
 export DATASET_PATH=gs://maxtext-dataset
 # Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you own, this bucket will store all the files generated by MaxText during a run
-export BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs
+export BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/gemm2-2b
 
 # We can also run finetuning by using the scanned converted checkpoint.
 # Note that scanned checkpoint helps with efficient finetuning

end_to_end/tpu/gemma3/4b/test_gemma3_to_mt.sh

@@ -75,7 +75,7 @@ fi
 # Non-Googlers please remember to point `DATASET_PATH` to the GCS bucket where you have your training data
 export DATASET_PATH=gs://maxtext-dataset
 # Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to a GCS bucket that you own, this bucket will store all the files generated by MaxText during a run
-export BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs
+export BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/gemma3-4b
 
 # We can also run finetuning by using the scanned converted checkpoint.
 # Note that scanned checkpoint helps with efficient finetuning

src/MaxText/utils/ckpt_scripts/llama_mistral_mixtral_orbax_to_hf.py

@@ -48,7 +48,7 @@
 from transformers import LlamaForCausalLM, MistralForCausalLM, AutoModelForCausalLM, AutoConfig
 
 from MaxText import checkpointing
-from MaxText import llama_or_mistral_ckpt
+from MaxText.utils.ckpt_scripts import llama_or_mistral_ckpt
 from MaxText import max_logging
 from MaxText import maxtext_utils
 from MaxText import pyconfig

src/MaxText/vllm_decode.py

@@ -2,10 +2,6 @@
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
-<<<<<<< HEAD
-=======
-
->>>>>>> c6a7412e (test)
 #  You may obtain a copy of the License at
 #
 #       https://www.apache.org/licenses/LICENSE-2.0
@@ -113,7 +109,9 @@ def main(argv: Sequence[str]) -> None:
   jax.config.update("jax_default_prng_impl", "unsafe_rbg")
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
   if "xla_tpu_spmd_rng_bit_generator_unsafe" not in os.environ.get("LIBTPU_INIT_ARGS", ""):
-    os.environ["LIBTPU_INIT_ARGS"] = os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
+    os.environ["LIBTPU_INIT_ARGS"] = (
+        os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
+    )
 
   config = pyconfig.initialize(argv)
   maxtext_model, mesh = model_creation_utils.create_nnx_model(config)