# relevant. All of the GPU test machines have CUDA installed via other means,
# so I am not sure how to verify that the dependencies themselves are valid for
# the moment.
if [[ "$TFCI_BAZEL_COMMON_ARGS" =~ gpu|cuda ]]; then
  echo "Checking to make sure tensorflow[and-cuda] is installable..."
  "$python" -m pip install "$(echo *.whl)[and-cuda]" $TFCI_PYTHON_VERIFY_PIP_INSTALL_ARGS

  # non-interactive shells as well.
  shopt -s expand_aliases
else
  echo '==TFCI==: Error: Cannot find path to grealpath or gstat'
  echo 'TF CI scripts require GNU core utilties to be installed. Please make'
  echo 'sure they are present on your system and try again.'
  exit 1
fi

# "TFCI_MACOS_BAZEL_TEST_DIR_PATH" specifies the directory that Bazel should use
# when running tests. Each test will be executed in a separate subdirectory

    TAG=$(head -n 1 "$KOKORO_PIPER_DIR/presubmit_request.txt" | cut -d" " -f2)
  else
    TAG="pr-${KOKORO_GITHUB_PULL_REQUEST_NUMBER}"
  fi
fi

# Build for both JAX and TF usage.  We do these in one place because they share
# almost all of the same cache layers
export DOCKER_BUILDKIT=1
for target in jax tf; do
  IMAGE="gcr.io/tensorflow-sigs/build-arm64:$target-$TAG"
  docker pull "$IMAGE" || true