From 39e8ee93b015372049c833bf45f49105026d9e8a Mon Sep 17 00:00:00 2001
From: Daniel Suo <dsuo@google.com>
Date: Fri, 14 Mar 2025 13:53:57 -0700
Subject: [PATCH] Add `experimental/serialize_executable.py` to `BUILD`.
PiperOrigin-RevId: 736975882
---
jax/BUILD | 17 +++++++++++++++++
jaxlib/jax.bzl | 1 +
2 files changed, 18 insertions(+)
diff --git a/jax/BUILD b/jax/BUILD
index 5e9cbd7a1..d6f100581 100644
--- a/jax/BUILD
+++ b/jax/BUILD
@@ -36,6 +36,7 @@ load(
"py_library_providing_imports_info",
"pytype_library",
"pytype_strict_library",
+ "serialize_executable_internal_users",
)
package(
@@ -127,6 +128,12 @@ package_group(
packages = mosaic_gpu_internal_users,
)
+package_group(
+ name = "serialize_executable_users",
+ includes = [":internal"],
+ packages = serialize_executable_internal_users,
+)
+
# JAX-private test utilities.
py_library(
# This build target is required in order to use private test utilities in jax._src.test_util,
@@ -618,6 +625,16 @@ pytype_strict_library(
] + py_deps("numpy"),
)
+pytype_strict_library(
+ name = "serialize_executable",
+ srcs = ["experimental/serialize_executable.py"],
+ visibility = [":serialize_executable_users"],
+ deps = [
+ ":jax",
+ "//jax/_src/lib",
+ ],
+)
+
pytype_strict_library(
name = "source_mapper",
srcs = glob(include = ["experimental/source_mapper/**/*.py"]),
diff --git a/jaxlib/jax.bzl b/jaxlib/jax.bzl
index c4ac8d00f..89f154599 100644
--- a/jaxlib/jax.bzl
+++ b/jaxlib/jax.bzl
@@ -48,6 +48,7 @@ pallas_gpu_internal_users = []
pallas_tpu_internal_users = []
pallas_fuser_users = []
mosaic_extension_deps = []
+serialize_executable_internal_users = []
jax_internal_export_back_compat_test_util_visibility = []
jax_internal_test_harnesses_visibility = []