diff --git a/WORKSPACE b/WORKSPACE
index dc83cf84..c9f36b97 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -1,7 +1,16 @@
 # add support for generating compile_commands
-load("//third_party/hedron_compile_commands:workspace.bzl", hedron_compile_commands_workspace = "repo", hedron_compile_commands_setup = "setup")
+load("//third_party/hedron_compile_commands:workspace.bzl", hedron_compile_commands_workspace = "repo")
 hedron_compile_commands_workspace()
+
+load("@hedron_compile_commands//:workspace_setup.bzl", "hedron_compile_commands_setup")
+load("@hedron_compile_commands//:workspace_setup_transitive.bzl", "hedron_compile_commands_setup_transitive")
+load("@hedron_compile_commands//:workspace_setup_transitive_transitive.bzl", "hedron_compile_commands_setup_transitive_transitive")
+load("@hedron_compile_commands//:workspace_setup_transitive_transitive_transitive.bzl", "hedron_compile_commands_setup_transitive_transitive_transitive")
+
 hedron_compile_commands_setup()
+hedron_compile_commands_setup_transitive()
+hedron_compile_commands_setup_transitive_transitive()
+hedron_compile_commands_setup_transitive_transitive_transitive()
 
 load("//third_party/jax:workspace.bzl", jax_workspace = "repo")
 jax_workspace()
@@ -45,6 +54,18 @@ llvm_configure(name = "llvm-project", targets = LLVM_TARGETS)
 load("@jax//third_party/flatbuffers:workspace.bzl", flatbuffers = "repo")
 flatbuffers()
 
+load("@xla//:workspace4.bzl", "xla_workspace4")
+load("@xla//:workspace3.bzl", "xla_workspace3")
+load("@xla//:workspace2.bzl", "xla_workspace2")
+load("@xla//:workspace1.bzl", "xla_workspace1")
+load("@xla//:workspace0.bzl", "xla_workspace0")
+
+xla_workspace4()
+xla_workspace3()
+xla_workspace2()
+xla_workspace1()
+xla_workspace0()
+
 load(
     "@tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl",
     "cuda_json_init_repository",
diff --git a/third_party/hedron_compile_commands/workspace.bzl b/third_party/hedron_compile_commands/workspace.bzl
index e7261047..bbb25a02 100644
--- a/third_party/hedron_compile_commands/workspace.bzl
+++ b/third_party/hedron_compile_commands/workspace.bzl
@@ -13,16 +13,3 @@ def repo():
         strip_prefix = "bazel-compile-commands-extractor-4f28899228fb3ad0126897876f147ca15026151e",
         # When you first run this tool, it'll recommend a sha256 hash to put here with a message like: "DEBUG: Rule 'hedron_compile_commands' indicated that a canonical reproducible form can be obtained by modifying arguments sha256 = ..."
     )
-
-def setup():
-    """Setups compile_commands generation."""
-
-    load("@hedron_compile_commands//:workspace_setup.bzl", "hedron_compile_commands_setup")
-    load("@hedron_compile_commands//:workspace_setup_transitive.bzl", "hedron_compile_commands_setup_transitive")
-    load("@hedron_compile_commands//:workspace_setup_transitive_transitive.bzl", "hedron_compile_commands_setup_transitive_transitive")
-    load("@hedron_compile_commands//:workspace_setup_transitive_transitive_transitive.bzl", "hedron_compile_commands_setup_transitive_transitive_transitive")
-
-    hedron_compile_commands_setup()
-    hedron_compile_commands_setup_transitive()
-    hedron_compile_commands_setup_transitive_transitive()
-    hedron_compile_commands_setup_transitive_transitive_transitive()
diff --git a/third_party/xla/workspace.bzl b/third_party/xla/workspace.bzl
index b029cc5a..d7feaaea 100644
--- a/third_party/xla/workspace.bzl
+++ b/third_party/xla/workspace.bzl
@@ -12,15 +12,3 @@ def repo():
         urls = ["https://github.com/wsmoses/xla/archive/{commit}.tar.gz".format(commit = XLA_COMMIT)],
         patch_cmds = XLA_PATCHES,
     )
-
-    load("@xla//:workspace4.bzl", "xla_workspace4")
-    load("@xla//:workspace3.bzl", "xla_workspace3")
-    load("@xla//:workspace2.bzl", "xla_workspace2")
-    load("@xla//:workspace1.bzl", "xla_workspace1")
-    load("@xla//:workspace0.bzl", "xla_workspace0")
-
-    xla_workspace4()
-    xla_workspace3()
-    xla_workspace2()
-    xla_workspace1()
-    xla_workspace0()