Add editable jax wheel target.

The set of editable wheels (`jaxlib`, `jax-cuda-plugin` and `jax-cuda-pjrt`) was used as dependencies in `requirements.in` file together with `:build_jaxlib=false` flag.

After [adding `jax` wheel dependencies](f5a4d1a) to the tests when `:build_jaxlib=false` is used, we need an editable `jax` wheel target as well to get the tests passing.

PiperOrigin-RevId: 740840736

Author: Google-ML-Automation

BUILD.bazel

@@ -72,35 +72,37 @@ py_binary(
     ],
 )
 
+WHEEL_SOURCE_FILES = [
+    ":transitive_py_data",
+    ":transitive_py_deps",
+    "//jax:py.typed",
+    "AUTHORS",
+    "LICENSE",
+    "README.md",
+    "pyproject.toml",
+    "setup.py",
+]
+
 jax_wheel(
     name = "jax_wheel",
     platform_independent = True,
-    source_files = [
-        ":transitive_py_data",
-        ":transitive_py_deps",
-        "//jax:py.typed",
-        "AUTHORS",
-        "LICENSE",
-        "README.md",
-        "pyproject.toml",
-        "setup.py",
-    ],
+    source_files = WHEEL_SOURCE_FILES,
+    wheel_binary = ":build_wheel",
+    wheel_name = "jax",
+)
+
+jax_wheel(
+    name = "jax_wheel_editable",
+    editable = True,
+    platform_independent = True,
+    source_files = WHEEL_SOURCE_FILES,
     wheel_binary = ":build_wheel",
     wheel_name = "jax",
 )
 
 jax_source_package(
     name = "jax_source_package",
-    source_files = [
-        ":transitive_py_data",
-        ":transitive_py_deps",
-        "//jax:py.typed",
-        "AUTHORS",
-        "LICENSE",
-        "README.md",
-        "pyproject.toml",
-        "setup.py",
-    ],
+    source_files = WHEEL_SOURCE_FILES,
     source_package_binary = ":build_wheel",
     source_package_name = "jax",
 )

build/build.py

@@ -68,10 +68,14 @@
 # rule as the default.
 WHEEL_BUILD_TARGET_DICT_NEW = {
     "jax": "//:jax_wheel",
+    "jax_editable": "//:jax_wheel_editable",
     "jax_source_package": "//:jax_source_package",
     "jaxlib": "//jaxlib/tools:jaxlib_wheel",
+    "jaxlib_editable": "//jaxlib/tools:jaxlib_wheel_editable",
     "jax-cuda-plugin": "//jaxlib/tools:jax_cuda_plugin_wheel",
+    "jax-cuda-plugin_editable": "//jaxlib/tools:jax_cuda_plugin_wheel_editable",
     "jax-cuda-pjrt": "//jaxlib/tools:jax_cuda_pjrt_wheel",
+    "jax-cuda-pjrt_editable": "//jaxlib/tools:jax_cuda_pjrt_wheel_editable",
     "jax-rocm-plugin": "//jaxlib/tools:jax_rocm_plugin_wheel",
     "jax-rocm-pjrt": "//jaxlib/tools:jax_rocm_pjrt_wheel",
 }
@@ -662,9 +666,12 @@ async def main():
       )
 
       # Append the build target to the Bazel command.
-      build_target = wheel_build_targets[wheel]
+      if args.use_new_wheel_build_rule and args.editable:
+        build_target = wheel_build_targets[wheel + "_editable"]
+      else:
+        build_target = wheel_build_targets[wheel]
       wheel_build_command.append(build_target)
-      if args.use_new_wheel_build_rule and wheel == "jax":
+      if args.use_new_wheel_build_rule and wheel == "jax" and not args.editable:
         wheel_build_command.append(wheel_build_targets["jax_source_package"])
 
       if not args.use_new_wheel_build_rule:

build_wheel.py

@@ -61,6 +61,11 @@
         "Whether to build the source package only. Optional."
     ),
 )
+parser.add_argument(
+    "--editable",
+    action="store_true",
+    help="Create an 'editable' jax build instead of a wheel.",
+)
 args = parser.parse_args()
 
 
@@ -90,7 +95,11 @@ def prepare_srcs(deps: list[str], srcs_dir: str) -> None:
   """
 
   for file in deps:
-    if not (file.startswith("bazel-out") or file.startswith("external")):
+    if not (
+        file.startswith("bazel-out")
+        or file.startswith("external")
+        or file.startswith("jaxlib")
+    ):
       copy_file(file, srcs_dir)
 
 
@@ -103,14 +112,18 @@ def prepare_srcs(deps: list[str], srcs_dir: str) -> None:
 try:
   os.makedirs(args.output_path, exist_ok=True)
   prepare_srcs(args.srcs, pathlib.Path(sources_path))
-  build_utils.build_wheel(
-      sources_path,
-      args.output_path,
-      package_name="jax",
-      git_hash=args.jaxlib_git_hash,
-      build_wheel_only=args.build_wheel_only,
-      build_source_package_only=args.build_source_package_only,
-  )
+  package_name = "jax"
+  if args.editable:
+    build_utils.build_editable(sources_path, args.output_path, package_name)
+  else:
+    build_utils.build_wheel(
+        sources_path,
+        args.output_path,
+        package_name,
+        git_hash=args.jaxlib_git_hash,
+        build_wheel_only=args.build_wheel_only,
+        build_source_package_only=args.build_source_package_only,
+    )
 finally:
   if tmpdir:
     tmpdir.cleanup()