Fix how xla is cloned at head for workflows

Author: nitins17

.github/workflows/bazel_cpu_rbe.yml

@@ -33,4 +33,4 @@ jobs:
       - name: Wait For Connection
         uses: ./actions/ci_connection/
       - name: Run Bazel CPU Tests
-        run: ./ci/run_bazel_test.sh "ci/envs/run_tests/bazel_cpu.env"
+        run: ./ci/run_bazel_test.sh "ci/envs/run_tests/bazel_cpu_rbe.env"

.github/workflows/bazel_gpu_non_rbe.yml

@@ -18,6 +18,7 @@ jobs:
     env:
       # GitHub actions run in Docker by defaut. Disable running the `setup_docker.sh` script.
       JAXCI_RUN_DOCKER_CONTAINER: 0
+      JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: 3.11
 
     steps:
@@ -32,6 +33,6 @@ jobs:
       - name: Wait For Connection
         uses: ./actions/ci_connection/
       - name: Run Bazel GPU tests locally
-        run: ./ci/run_bazel_test.sh "ci/envs/run_tests/bazel_gpu_local"
+        run: ./ci/run_bazel_test.sh "ci/envs/run_tests/bazel_gpu_non_rbe.env"
       - name: Test if step runs
         run: echo "This step ran!"

.github/workflows/bazel_gpu_rbe.yml

@@ -18,6 +18,7 @@ jobs:
     env:
       # Do not run Docker container for Linux runners. Linux runners already run in a Docker container.
       JAXCI_RUN_DOCKER_CONTAINER: 0
+      JAXCI_CLONE_MAIN_XLA: 1
       JAXCI_HERMETIC_PYTHON_VERSION: 3.12
 
     steps:

ci/envs/run_tests/bazel_gpu_non_rbe.env

@@ -21,14 +21,3 @@ export JAXCI_RUN_BAZEL_TEST_GPU_LOCAL=1
 # Only Linux x86 runs local GPU tests at the moment.
 export JAXCI_DOCKER_IMAGE="gcr.io/tensorflow-testing/nosla-cuda${JAXCI_DOCKER_CUDA_VERSION}-cudnn9.1-ubuntu20.04-manylinux2014-multipython"
 export JAXCI_DOCKER_ARGS="--shm-size=16g --gpus all"
-
-export NCCL_DEBUG=WARN
-
-# Disable x64 mode
-export JAX_ENABLE_X64=0
-
-# Clone XLA at HEAD.
-export JAXCI_CLONE_MAIN_XLA=1
-
-# Set per device memory limit to 20.5 GB
-export TF_PER_DEVICE_MEMORY_LIMIT_MB=20480

ci/envs/run_tests/bazel_gpu_rbe.env

@@ -26,8 +26,5 @@ export JAXCI_DOCKER_ARGS="--gpus all"
 export JAX_CUDA_VERSION=12
 export JAX_CUDNN_VERSION=9.1
 
-# Disable x64 mode
-export JAX_ENABLE_X64=0
-
 # Clone XLA at HEAD.
 export JAXCI_CLONE_MAIN_XLA=1