Migrate jaxlib to use a single common .so file for all C++ dependencies.

The idea is to move all of the jaxlib contents into a single .so file,
and have all of the other Python extensions be tiny stubs that reexport
part of the larger .so file.

This has two main benefits:
* it reduces the size of the jaxlib wheel, by about 70-80MB when
  installed. The benefit of the change is that it avoid duplication
  between the MLIR CAPI code and the copy of MLIR in XLA.
* it gives us flexibility to split and merge Python extensions as we see
  fit.

Issue #11225

PiperOrigin-RevId: 744855997

Author: hawkinsp

.bazelrc

@@ -98,6 +98,7 @@ build:windows --incompatible_strict_action_env=true
 # #############################################################################
 build:nonccl --define=no_nccl_support=true
 
+build --repo_env USE_PYWRAP_RULES=1
 build:posix --copt=-fvisibility=hidden
 build:posix --copt=-Wno-sign-compare
 build:posix --cxxopt=-std=c++17

jaxlib/BUILD

@@ -20,6 +20,12 @@ load(
     "py_library_providing_imports_info",
     "pytype_library",
 )
+load(
+    "//jaxlib:pywrap.bzl",
+    "nanobind_pywrap_extension",
+    "pywrap_binaries",
+    "pywrap_library",
+)
 load("//jaxlib:symlink_files.bzl", "symlink_files")
 
 licenses(["notice"])
@@ -51,6 +57,7 @@ py_library_providing_imports_info(
     lib_rule = pytype_library,
     deps = [
         ":cpu_feature_guard",
+        ":jax",
         ":utils",
         "//jaxlib/cpu:_lapack",
         "//jaxlib/mlir",
@@ -98,6 +105,44 @@ exports_files([
     "setup.py",
 ])
 
+pywrap_library(
+    name = "jax",
+    common_lib_def_files_or_filters = {
+        "jaxlib/jax_common": "jax_common.json",
+    },
+    common_lib_version_scripts = {
+        "jaxlib/jax_common": select({
+            "@bazel_tools//src/conditions:windows": None,
+            "@bazel_tools//src/conditions:darwin": "libjax_common_darwin.lds",
+            "//conditions:default": "libjax_common.lds",
+        }),
+    },
+    deps = [
+        ":utils",
+        "//jaxlib/mlir/_mlir_libs:_chlo",
+        "//jaxlib/mlir/_mlir_libs:_mlir",
+        "//jaxlib/mlir/_mlir_libs:_mlirDialectsGPU",
+        "//jaxlib/mlir/_mlir_libs:_mlirDialectsLLVM",
+        "//jaxlib/mlir/_mlir_libs:_mlirDialectsNVGPU",
+        "//jaxlib/mlir/_mlir_libs:_mlirDialectsSparseTensor",
+        "//jaxlib/mlir/_mlir_libs:_mlirGPUPasses",
+        "//jaxlib/mlir/_mlir_libs:_mlirHlo",
+        "//jaxlib/mlir/_mlir_libs:_mlirSparseTensorPasses",
+        "//jaxlib/mlir/_mlir_libs:_mosaic_gpu_ext",
+        "//jaxlib/mlir/_mlir_libs:_sdy",
+        "//jaxlib/mlir/_mlir_libs:_stablehlo",
+        "//jaxlib/mlir/_mlir_libs:_tpu_ext",
+        "//jaxlib/mlir/_mlir_libs:_triton_ext",
+        "//jaxlib/mlir/_mlir_libs:register_jax_dialects",
+        "//jaxlib/xla:xla_extension",
+    ],
+)
+
+pywrap_binaries(
+    name = "jaxlib_binaries",
+    dep = ":jax",
+)
+
 cc_library(
     name = "absl_status_casters",
     hdrs = ["absl_status_casters.h"],
@@ -170,10 +215,9 @@ nanobind_extension(
     ],
 )
 
-nanobind_extension(
+nanobind_pywrap_extension(
     name = "utils",
     srcs = ["utils.cc"],
-    module_name = "utils",
     deps = [
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",

jaxlib/jax_common.json

@@ -0,0 +1,8 @@
+{
+  "global": [
+    "Wrapped_PyInit_*"
+  ],
+  "local": [
+    "*"
+  ]
+}

jaxlib/libjax_common.lds

@@ -0,0 +1,7 @@
+{
+  global:
+    Wrapped_PyInit_*;
+
+  local:
+    *;
+};

jaxlib/libjax_common_darwin.lds

@@ -0,0 +1 @@
+*Wrapped_PyInit_*