update cuda-pjrt env to fix sourcing default envs and remove redundant settings

Author: nitins17

ci/envs/build_artifacts/jax-cuda-pjrt.env

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 # Inherit default JAXCI environment variables.
-source ci/envs/default
+source ci/envs/default.env
 
 # Enable jax-cuda-pjrt build.
 export JAXCI_BUILD_PJRT="1"
@@ -35,13 +35,4 @@ fi
 # Linux Aarch64 specifc settings
 if [[ $os == "linux" ]] && [[ $arch == "aarch64" ]]; then
   export JAXCI_DOCKER_IMAGE="us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/build-arm64:tf-2-18-multi-python"
-fi
-
-# When building artifacts for running tests, we need to disable x64 mode and
-# clone XLA at HEAD.
-if [[ $JAXCI_SETUP_TEST_ENVIRONMENT == 1 ]]; then
-  # Disable x64 mode
-  export JAX_ENABLE_X64=0
-  # Clone XLA at HEAD.
-  export JAXCI_CLONE_MAIN_XLA=1
-fi
+fi

ci/utilities/setup_jaxci_envs.sh

@@ -27,6 +27,7 @@ if [[ -z "$1" ]]; then
 fi
 
 # -e: abort script if one command fails
+# -u: error if undefined variable used
 # -x: log all commands
 # -o history: record shell history
 # -o allexport: export all functions and variables to be available to subscripts