Rename arg in build script to be more clear.

The flag means skips GPU plugin extension in jaxlib.

PiperOrigin-RevId: 627203738

Author: Jieying Luo

build/build.py

@@ -674,7 +674,7 @@ def main():
       f"--jaxlib_git_hash={get_githash()}",
       f"--cpu={wheel_cpu}"])
     if args.build_gpu_plugin:
-      command.append("--include_gpu_plugin_extension")
+      command.append("--skip_gpu_kernels")
     if args.editable:
       command += ["--editable"]
     print(" ".join(command))

jaxlib/tools/build_wheel.py

@@ -57,11 +57,11 @@
     help="Create an 'editable' jaxlib build instead of a wheel.",
 )
 parser.add_argument(
-    "--include_gpu_plugin_extension",
-    # args.include_gpu_plugin_extension is True when
-    # --include_gpu_plugin_extension is in the command
+    "--skip_gpu_kernels",
+    # args.skip_gpu_kernels is True when
+    # --skip_gpu_kernels is in the command
     action="store_true",
-    help="Whether to include gpu plugin extension.",
+    help="Whether to skip gpu kernels in jaxlib.",
 )
 args = parser.parse_args()
 
@@ -169,7 +169,7 @@ def write_setup_cfg(sources_path, cpu):
     )
 
 
-def prepare_wheel(sources_path: pathlib.Path, *, cpu, include_gpu_plugin_extension):
+def prepare_wheel(sources_path: pathlib.Path, *, cpu, skip_gpu_kernels):
   """Assembles a source tree for the wheel in `sources_path`."""
   copy_runfiles = functools.partial(build_utils.copy_file, runfiles=r)
 
@@ -222,7 +222,7 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu, include_gpu_plugin_extensi
       ],
   )
 
-  if exists(f"__main__/jaxlib/cuda/_solver.{pyext}") and not include_gpu_plugin_extension:
+  if exists(f"__main__/jaxlib/cuda/_solver.{pyext}") and not skip_gpu_kernels:
     copy_runfiles(
         dst_dir=jaxlib_dir / "cuda" / "nvvm" / "libdevice",
         src_files=["local_config_cuda/cuda/cuda/nvvm/libdevice/libdevice.10.bc"],
@@ -413,7 +413,7 @@ def if_has_mosaic_gpu(extras):
   prepare_wheel(
       pathlib.Path(sources_path),
       cpu=args.cpu,
-      include_gpu_plugin_extension=args.include_gpu_plugin_extension,
+      skip_gpu_kernels=args.skip_gpu_kernels,
   )
   package_name = "jaxlib"
   if args.editable: