test

Author: SurbhiJainUSC

src/MaxText/vllm_decode.py

@@ -2,6 +2,10 @@
 #
 #  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
@@ -109,9 +113,7 @@ 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)