Change the definition of the config to run bazel test for cuda plug…

…in to match `//jax:build_jaxlib`.

When build_cuda_plugin_from_source is true, it will build cuda plugin from source, and it is used for the case of `bazel test` without preinstall jax cuda packages.

PiperOrigin-RevId: 582892405

jax/BUILD

@@ -52,18 +52,17 @@ config_setting(
     },
 )
 
-# When `build_cuda_plugin_for_tests` is true, it assumes `bazel build` for the cuda plugin will run
-# before `bazel test` for cuda plugin tests. Set it to false for the case of running `bazel test`
-# without `bazel build` for the cuda plugin.
+# When `build_cuda_plugin_from_source` is true, it assumes running `bazel test` without preinstalled
+# cuda plugin.
 bool_flag(
-    name = "build_cuda_plugin_for_tests",
-    build_setting_default = True,
+    name = "build_cuda_plugin_from_source",
+    build_setting_default = False,
 )
 
 config_setting(
-    name = "enable_cuda_plugin_build_for_tests",
+    name = "enable_build_cuda_plugin_from_source",
     flag_values = {
-        ":build_cuda_plugin_for_tests": "True",
+        ":build_cuda_plugin_from_source": "True",
     },
 )
 

jaxlib/jax.bzl

@@ -202,8 +202,8 @@ def jax_test(
                 "//jax:enable_jaxlib_build": ["//jaxlib/cuda:gpu_only_test_deps"],
                 "//conditions:default": [],
             }) + select({
-                "//jax:enable_cuda_plugin_build_for_tests": [],
-                "//conditions:default": ["//jax_plugins:gpu_plugin_only_test_deps"],
+                "//jax:enable_build_cuda_plugin_from_source": ["//jax_plugins:gpu_plugin_only_test_deps"],
+                "//conditions:default": [],
             }),
             shard_count = test_shards,
             tags = test_tags,