Move the XLA commit out of the top-level JAX WORKSPACE file and into a separate .bzl file.

No functional changes intended.

PiperOrigin-RevId: 556906943

WORKSPACE

@@ -1,31 +1,6 @@
-load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
-
-# To update XLA to a new revision,
-# a) update URL and strip_prefix to the new git commit hash
-# b) get the sha256 hash of the commit by running:
-#    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
-#    and update the sha256 with the result.
-http_archive(
-    name = "xla",
-    sha256 = "45b1d9b0a25130b5805184278e0e3342ed95b7638c1a1fc6fd663188e6784401",
-    strip_prefix = "xla-9f26b9390f5a5c565a13925731de749be8a760be",
-    urls = [
-        "https://github.com/openxla/xla/archive/9f26b9390f5a5c565a13925731de749be8a760be.tar.gz",
-    ],
-)
-
-# For development, one often wants to make changes to the TF repository as well
-# as the JAX repository. You can override the pinned repository above with a
-# local checkout by either:
-# a) overriding the TF repository on the build.py command line by passing a flag
-#    like:
-#    python build/build.py --bazel_options=--override_repository=xla=/path/to/xla
-#    or
-# b) by commenting out the http_archive above and uncommenting the following:
-# local_repository(
-#    name = "xla",
-#    path = "/path/to/xla",
-# )
+# The XLA commit is determined by third_party/xla/workspace.bzl.
+load("//third_party/xla:workspace.bzl", jax_xla_workspace = "repo")
+jax_xla_workspace()
 
 load("//third_party/ducc:workspace.bzl", ducc = "repo")
 ducc()

third_party/xla/workspace.bzl

@@ -0,0 +1,47 @@
+# Copyright 2023 The JAX Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+
+# To update XLA to a new revision,
+# a) update XLA_COMMIT to the new git commit hash
+# b) get the sha256 hash of the commit by running:
+#    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
+#    and update XLA_SHA256 with the result.
+
+XLA_COMMIT = "9f26b9390f5a5c565a13925731de749be8a760be"
+XLA_SHA256 = "45b1d9b0a25130b5805184278e0e3342ed95b7638c1a1fc6fd663188e6784401"
+
+def repo():
+    http_archive(
+        name = "xla",
+        sha256 = XLA_SHA256,
+        strip_prefix = "xla-{commit}".format(commit = XLA_COMMIT),
+        urls = [
+            "https://github.com/openxla/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT),
+        ],
+    )
+
+    # For development, one often wants to make changes to the TF repository as well
+    # as the JAX repository. You can override the pinned repository above with a
+    # local checkout by either:
+    # a) overriding the TF repository on the build.py command line by passing a flag
+    #    like:
+    #    python build/build.py --bazel_options=--override_repository=xla=/path/to/xla
+    #    or
+    # b) by commenting out the http_archive above and uncommenting the following:
+    # local_repository(
+    #    name = "xla",
+    #    path = "/path/to/xla",
+    # )