[JAX] Remove flag jax_cpu_enable_gloo_collectives.

This flag is (a) on by default and (b) replaced by jax_cpu_collectives_implementation.

PiperOrigin-RevId: 808183075

Author: hawkinsp

src/MaxText/standalone_checkpointer.py

@@ -117,7 +117,6 @@ def add_entropy_to_checkpoint(state):
 
 
 def main(argv: Sequence[str]) -> None:
-  jax.config.update("jax_cpu_enable_gloo_collectives", True)
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
   config = pyconfig.initialize(argv)
   validate_train_config(config)

src/MaxText/standalone_dataloader.py

@@ -65,7 +65,6 @@ def data_load_loop(config, state=None):
 
 
 def main(argv: Sequence[str]) -> None:
-  jax.config.update("jax_cpu_enable_gloo_collectives", True)
   os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
   config = pyconfig.initialize(argv)
   validate_train_config(config)