diff --git a/MODULE.bazel b/MODULE.bazel
index 940ed0be..ac03005a 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -1,36 +1,7 @@
-JDK_VERSIONS = [
-    "11",
-    "15",
-    "16",
-    "17",
-]
-
-PLATFORMS = [
-    "linux",
-    "macos",
-    "macos_aarch64",
-    "win",
-]
-
-# Remote JDK repos for those Linux platforms are only defined for JDK 11.
-EXTRA_REMOTE_JDK11_REPOS = [
-    "remotejdk11_linux_aarch64",
-    "remotejdk11_linux_ppc64le",
-    "remotejdk11_linux_s390x",
-]
-
-REMOTE_JDK_REPOS = [("remotejdk" + version + "_" + platform) for version in JDK_VERSIONS for platform in PLATFORMS] + EXTRA_REMOTE_JDK11_REPOS
-
-REMOTE_JAVA_TOOLCHAINS = [("@" + name + "_toolchain_config_repo//:toolchain") for name in REMOTE_JDK_REPOS]
-
 module(
     name = "rules_java",
     compatibility_level = 1,
-    toolchains_to_register = [
-        "//toolchains:all",
-        "@local_jdk//:runtime_toolchain_definition",
-    ] + REMOTE_JAVA_TOOLCHAINS,
-    version = "4.0.0",
+    version = "5.1.0",
 )
 
 bazel_dep(name = "platforms", version = "0.0.4")
@@ -39,6 +10,8 @@ bazel_dep(name = "rules_cc", version = "0.0.1")
 # rules_proto is required by @remote_java_tools, which is loaded via module extension.
 bazel_dep(name = "rules_proto", version = "4.0.0")
 
+register_toolchains("//toolchains:all")
+
 toolchains = use_extension("//java:extensions.bzl", "toolchains")
 
 # Declare remote java tools repos
@@ -53,12 +26,39 @@ use_repo(toolchains, "remote_java_tools_darwin")
 # Declare local jdk repo
 use_repo(toolchains, "local_jdk")
 
+register_toolchains("@local_jdk//:runtime_toolchain_definition")
+
 # Declare all remote jdk toolchain config repos
+JDK_VERSIONS = [
+    "11",
+    "15",
+    "16",
+    "17",
+]
+
+PLATFORMS = [
+    "linux",
+    "macos",
+    "macos_aarch64",
+    "win",
+]
+
+# Remote JDK repos for those Linux platforms are only defined for JDK 11.
+EXTRA_REMOTE_JDK11_REPOS = [
+    "remotejdk11_linux_aarch64",
+    "remotejdk11_linux_ppc64le",
+    "remotejdk11_linux_s390x",
+]
+
+REMOTE_JDK_REPOS = [("remotejdk" + version + "_" + platform) for version in JDK_VERSIONS for platform in PLATFORMS] + EXTRA_REMOTE_JDK11_REPOS
+
 [use_repo(
     toolchains,
     repo + "_toolchain_config_repo",
 ) for repo in REMOTE_JDK_REPOS]
 
+[register_toolchains("@" + name + "_toolchain_config_repo//:toolchain") for name in REMOTE_JDK_REPOS]
+
 # Dev dependencies
 bazel_dep(name = "rules_pkg", dev_dependency = True, version = "0.5.1")
 bazel_dep(name = "bazel_skylib", dev_dependency = True, version = "1.2.0")