Update XLA. (#1702)

Add support for building a CPU-only jaxlib with a CUDA-enabled toolchain.

WORKSPACE

@@ -23,10 +23,10 @@ http_archive(
 #    and update the sha256 with the result.
 http_archive(
     name = "org_tensorflow",
-    sha256 = "cf32fabf086a12f52d1e3dd87eac9b44d44f76713030c45e22d1913ed1edb4d0",
-    strip_prefix = "tensorflow-d55375021a36e404015abb3439e631f9fac59838",
+    sha256 = "f0be98a469fcefab7a48721db63ea28fab24c08fd5171214978f88f02d0f07ad",
+    strip_prefix = "tensorflow-2bc7129de76dc6fdaab5619aaecf7a6bbfa17104",
     urls = [
-        "https://github.com/tensorflow/tensorflow/archive/d55375021a36e404015abb3439e631f9fac59838.tar.gz",
+        "https://github.com/tensorflow/tensorflow/archive/2bc7129de76dc6fdaab5619aaecf7a6bbfa17104.tar.gz",
     ],
 )
 

build/build.py

@@ -280,6 +280,13 @@ def main():
       parser,
       "enable_cuda",
       help_str="Should we build with CUDA enabled? Requires CUDA and CuDNN.")
+  add_boolean_argument(
+      parser,
+      "include_gpu_backend_if_cuda_enabled",
+      default=True,
+      help_str="If CUDA is enabled, should we build the GPU backend? This is "
+               "mostly useful to build a CPU-only build with a CUDA-enabled "
+               "toolchain.")
   parser.add_argument(
       "--cuda_path",
       default=None,
@@ -334,6 +341,8 @@ def main():
     config_args += ["--config=mkl_open_source_only"]
   if args.enable_cuda:
     config_args += ["--config=cuda"]
+    if args.include_gpu_backend_if_cuda_enabled:
+      config_args += ["--define=xla_python_enable_gpu=true"]
   shell(
     [bazel_path] + args.bazel_startup_options +
     ["run", "--verbose_failures=true"] + config_args +

build/build_wheel_docker_entrypoint.sh

@@ -49,7 +49,8 @@ case $2 in
     PLAT_NAME="linux_x86_64"
     ;;
   nocuda)
-    python build.py --bazel_startup_options="--output_user_root=/build/root"
+    python build.py --enable_cuda --include_gpu_backend_if_cuda_enabled=false \
+      --bazel_startup_options="--output_user_root=/build/root"
     PLAT_NAME="manylinux2010_x86_64"
     ;;
   *)