From 3c82c787170af4a0da4d371c892d3574d5cdcc85 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 6 Feb 2025 15:46:51 -0800
Subject: [PATCH] Create JAX wheel build target.

This change introduces a uniform way of building the artifacts and controlling the filename version suffixes (see the changes for jaxlib, CUDA and PJRT plugins in https://github.com/jax-ml/jax/pull/25126)

Previously JAX wheel was built via `python3 -m build` command. The resulting wheel contained the python packages files in `jax` folder (e.g. the files in the subdirs that have `__init__.py` file).

You can still build the JAX wheel with `python3 -m build` command.

Bazel command example for building nightly JAX wheel:

```
bazel build :jax_wheel \
  --config=ci_linux_x86_64 \
  --repo_env=HERMETIC_PYTHON_VERSION=3.10 \
  --repo_env=ML_WHEEL_TYPE=custom \
  --repo_env=ML_WHEEL_BUILD_DATE=20250211 \
  --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD)
```

Resulting wheel:
```
bazel-bin/dist/jax-0.5.1.dev20250211+d4f1f2278-py3-none-any.whl
```

PiperOrigin-RevId: 724102315
---
 third_party/py/python_wheel.bzl | 95 +++++++++++++++++++++------------
 1 file changed, 62 insertions(+), 33 deletions(-)

diff --git a/third_party/py/python_wheel.bzl b/third_party/py/python_wheel.bzl
index 922c2d392..89153e742 100644
--- a/third_party/py/python_wheel.bzl
+++ b/third_party/py/python_wheel.bzl
@@ -1,36 +1,4 @@
-"""
-Python wheel repository rules.
-
-The calculated wheel version suffix depends on the wheel type:
-- nightly: .dev{build_date}
-- release: ({custom_version_suffix})?
-- custom: .dev{build_date}(+{git_hash})?({custom_version_suffix})?
-- snapshot (default): -0
-
-The following environment variables can be set:
-{wheel_type}: ML_WHEEL_TYPE
-{build_date}: ML_WHEEL_BUILD_DATE (should be YYYYMMDD or YYYY-MM-DD)
-{git_hash}: ML_WHEEL_GIT_HASH
-{custom_version_suffix}: ML_WHEEL_VERSION_SUFFIX
-
-Examples:
-1. nightly wheel version: 2.19.0.dev20250107
-   Env vars passed to Bazel command: --repo_env=ML_WHEEL_TYPE=nightly
-                                     --repo_env=ML_WHEEL_BUILD_DATE=20250107
-2. release wheel version: 2.19.0
-   Env vars passed to Bazel command: --repo_env=ML_WHEEL_TYPE=release
-3. release candidate wheel version: 2.19.0-rc1
-   Env vars passed to Bazel command: --repo_env=ML_WHEEL_TYPE=release
-                                     --repo_env=ML_WHEEL_VERSION_SUFFIX=-rc1
-4. custom wheel version: 2.19.0.dev20250107+cbe478fc5-custom
-   Env vars passed to Bazel command: --repo_env=ML_WHEEL_TYPE=custom
-                                     --repo_env=ML_WHEEL_BUILD_DATE=$(git show -s --format=%as HEAD)
-                                     --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD)
-                                     --repo_env=ML_WHEEL_VERSION_SUFFIX=-custom
-5. snapshot wheel version: 2.19.0-0
-   Env vars passed to Bazel command: --repo_env=ML_WHEEL_TYPE=snapshot
-
-"""
+""" Repository and build rules for Python wheels packaging utilities. """
 
 def _get_host_environ(repository_ctx, name, default_value = None):
     """Returns the value of an environment variable on the host platform.
@@ -127,3 +95,64 @@ python_wheel_version_suffix_repository = repository_rule(
     implementation = _python_wheel_version_suffix_repository_impl,
     environ = _ENVIRONS,
 )
+
+""" Repository rule for storing Python wheel filename version suffix.
+
+The calculated wheel version suffix depends on the wheel type:
+- nightly: .dev{build_date}
+- release: ({custom_version_suffix})?
+- custom: .dev{build_date}(+{git_hash})?({custom_version_suffix})?
+- snapshot (default): -0
+
+The following environment variables can be set:
+{wheel_type}: ML_WHEEL_TYPE
+{build_date}: ML_WHEEL_BUILD_DATE (should be YYYYMMDD or YYYY-MM-DD)
+{git_hash}: ML_WHEEL_GIT_HASH
+{custom_version_suffix}: ML_WHEEL_VERSION_SUFFIX
+
+Examples:
+1. nightly wheel version: 2.19.0.dev20250107
+   Env vars passed to Bazel command: --repo_env=ML_WHEEL_TYPE=nightly
+                                     --repo_env=ML_WHEEL_BUILD_DATE=20250107
+2. release wheel version: 2.19.0
+   Env vars passed to Bazel command: --repo_env=ML_WHEEL_TYPE=release
+3. release candidate wheel version: 2.19.0-rc1
+   Env vars passed to Bazel command: --repo_env=ML_WHEEL_TYPE=release
+                                     --repo_env=ML_WHEEL_VERSION_SUFFIX=-rc1
+4. custom wheel version: 2.19.0.dev20250107+cbe478fc5-custom
+   Env vars passed to Bazel command: --repo_env=ML_WHEEL_TYPE=custom
+                                     --repo_env=ML_WHEEL_BUILD_DATE=$(git show -s --format=%as HEAD)
+                                     --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD)
+                                     --repo_env=ML_WHEEL_VERSION_SUFFIX=-custom
+5. snapshot wheel version: 2.19.0-0
+   Env vars passed to Bazel command: --repo_env=ML_WHEEL_TYPE=snapshot
+
+"""  # buildifier: disable=no-effect
+
+def _transitive_py_deps_impl(ctx):
+    outputs = depset(
+        [],
+        transitive = [dep[PyInfo].transitive_sources for dep in ctx.attr.deps],
+    )
+
+    return DefaultInfo(files = outputs)
+
+_transitive_py_deps = rule(
+    attrs = {
+        "deps": attr.label_list(
+            allow_files = True,
+            providers = [PyInfo],
+        ),
+    },
+    implementation = _transitive_py_deps_impl,
+)
+
+def transitive_py_deps(name, deps = []):
+    _transitive_py_deps(name = name + "_gather", deps = deps)
+    native.filegroup(name = name, srcs = [":" + name + "_gather"])
+
+"""Collects python files that a target depends on.
+
+It traverses dependencies of provided targets, collect their direct and 
+transitive python deps and then return a list of paths to files.
+"""  # buildifier: disable=no-effect