Repository: jax-ml/jax (python)

Issue:
`jnp.linalg.eigh` accesses full matrix rather than lower part
The behavior of `jnp.linalg.eigh` differs from the behavior of `np.linalg.eigh`:
```python
import numpy as np
import jax.numpy as jnp

x = np.array([[1, 0, 0],
              [2, 1, 0],
              [3, 2, 1]], dtype=float)

print(np.linalg.eigh(x)[0])
# [-2.         -0.70156212  5.70156212]

print(jnp.linalg.eigh(x)[0])
# [-0.49999994  0.14921923  3.3507814 ]

print(jnp.linalg.eigh(x, symmetrize_input=False)[0])  # This recovers the same results as NumPy.
# [-2.0000002 -0.7015623  5.701563 ]
```

This is confusing, and we should probably change the default value of `symmetrize_input`, unless there is good reason not to.

Related considerations:
- `symmetrize_input` should be documented in `jnp.linalg.eigh` using `lax_description`
- `jnp.linalg.eigvalsh` should probably support this argument as well.

Pull request description:
jnp.linalg: add symmetrize_input argument & docs
Fixes #9473.

At this point I think a change in default behavior is probably not in the cards, but now the functions document the difference in APIs with respect to NumPy, and allow for control over this via the `symmetrize_input` argument.

Patch:
diff --git a/.bazelrc b/.bazelrc
index 0c359e039..422363644 100644
--- a/.bazelrc
+++ b/.bazelrc
@@ -98,7 +98,6 @@ 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
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
index 23f2a58b0..146bbbda0 100644
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -72,8 +72,8 @@ def _symmetrize(x: Array) -> Array: return (x + _H(x)) / 2
 
 
 @export
-@partial(jit, static_argnames=['upper'])
-def cholesky(a: ArrayLike, *, upper: bool = False) -> Array:
+@partial(jit, static_argnames=['upper', 'symmetrize_input'])
+def cholesky(a: ArrayLike, *, upper: bool = False, symmetrize_input: bool = True) -> Array:
   """Compute the Cholesky decomposition of a matrix.
 
   JAX implementation of :func:`numpy.linalg.cholesky`.
@@ -98,6 +98,10 @@ def cholesky(a: ArrayLike, *, upper: bool = False) -> Array:
       Must have shape ``(..., N, N)``.
     upper: if True, compute the upper Cholesky decomposition `U`. if False
       (default), compute the lower Cholesky decomposition `L`.
+    symmetrize_input: if True (default) then input is symmetrized, which leads
+      to better behavior under automatic differentiation. Note that when this
+      is set to True, both the upper and lower triangles of the input will
+      be used in computing the decomposition.
 
   Returns:
     array of shape ``(..., N, N)`` representing the Cholesky decomposition
@@ -135,7 +139,7 @@ def cholesky(a: ArrayLike, *, upper: bool = False) -> Array:
   """
   a = ensure_arraylike("jnp.linalg.cholesky", a)
   a, = promote_dtypes_inexact(a)
-  L = lax_linalg.cholesky(a)
+  L = lax_linalg.cholesky(a, symmetrize_input=symmetrize_input)
   return L.mT.conj() if upper else L
 
 
@@ -821,7 +825,9 @@ def eigh(a: ArrayLike, UPLO: str | None = None,
     UPLO: specifies whether the calculation is done with the lower triangular
       part of ``a`` (``'L'``, default) or the upper triangular part (``'U'``).
     symmetrize_input: if True (default) then input is symmetrized, which leads
-      to better behavior under automatic differentiation.
+      to better behavior under automatic differentiation. Note that when this
+      is set to True, both the upper and lower triangles of the input will
+      be used in computing the decomposition.
 
   Returns:
     A namedtuple ``(eigenvalues, eigenvectors)`` where
@@ -863,8 +869,9 @@ def eigh(a: ArrayLike, UPLO: str | None = None,
 
 
 @export
-@partial(jit, static_argnames=('UPLO',))
-def eigvalsh(a: ArrayLike, UPLO: str | None = 'L') -> Array:
+@partial(jit, static_argnames=('UPLO', 'symmetrize_input'))
+def eigvalsh(a: ArrayLike, UPLO: str | None = 'L', *,
+             symmetrize_input: bool = True) -> Array:
   """
   Compute the eigenvalues of a Hermitian matrix.
 
@@ -875,6 +882,10 @@ def eigvalsh(a: ArrayLike, UPLO: str | None = 'L') -> Array:
       or symmetric (if real) matrix.
     UPLO: specifies whether the calculation is done with the lower triangular
       part of ``a`` (``'L'``, default) or the upper triangular part (``'U'``).
+    symmetrize_input: if True (default) then input is symmetrized, which leads
+      to better behavior under automatic differentiation. Note that when this
+      is set to True, both the upper and lower triangles of the input will
+      be used in computing the decomposition.
 
   Returns:
     An array of shape ``(..., M)`` containing the eigenvalues, sorted in
@@ -894,7 +905,7 @@ def eigvalsh(a: ArrayLike, UPLO: str | None = 'L') -> Array:
   """
   a = ensure_arraylike("jnp.linalg.eigvalsh", a)
   a, = promote_dtypes_inexact(a)
-  w, _ = eigh(a, UPLO)
+  w, _ = eigh(a, UPLO, symmetrize_input=symmetrize_input)
   return w
 
 
diff --git a/jaxlib/BUILD b/jaxlib/BUILD
index d195bda41..c8114b488 100644
--- a/jaxlib/BUILD
+++ b/jaxlib/BUILD
@@ -20,12 +20,6 @@ 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"])
@@ -57,7 +51,6 @@ py_library_providing_imports_info(
     lib_rule = pytype_library,
     deps = [
         ":cpu_feature_guard",
-        ":jax",
         ":utils",
         "//jaxlib/cpu:_lapack",
         "//jaxlib/mlir",
@@ -105,44 +98,6 @@ 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"],
@@ -215,9 +170,10 @@ nanobind_extension(
     ],
 )
 
-nanobind_pywrap_extension(
+nanobind_extension(
     name = "utils",
     srcs = ["utils.cc"],
+    module_name = "utils",
     deps = [
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
diff --git a/jaxlib/jax_common.json b/jaxlib/jax_common.json
deleted file mode 100644
index 61a2c9313..000000000
--- a/jaxlib/jax_common.json
+++ /dev/null
@@ -1,8 +0,0 @@
-{
-  "global": [
-    "Wrapped_PyInit_*"
-  ],
-  "local": [
-    "*"
-  ]
-}
diff --git a/jaxlib/libjax_common.lds b/jaxlib/libjax_common.lds
deleted file mode 100644
index 6130415a8..000000000
--- a/jaxlib/libjax_common.lds
+++ /dev/null
@@ -1,7 +0,0 @@
-{
-  global:
-    Wrapped_PyInit_*;
-
-  local:
-    *;
-};
diff --git a/jaxlib/libjax_common_darwin.lds b/jaxlib/libjax_common_darwin.lds
deleted file mode 100644
index aed9a1d75..000000000
--- a/jaxlib/libjax_common_darwin.lds
+++ /dev/null
@@ -1,1 +0,0 @@
-*Wrapped_PyInit_*
diff --git a/jaxlib/mlir/_mlir_libs/BUILD.bazel b/jaxlib/mlir/_mlir_libs/BUILD.bazel
index 6599e5069..fb94837cf 100644
--- a/jaxlib/mlir/_mlir_libs/BUILD.bazel
+++ b/jaxlib/mlir/_mlir_libs/BUILD.bazel
@@ -15,9 +15,10 @@
 load(
     "//jaxlib:jax.bzl",
     "if_windows",
+    "nanobind_extension",
+    "py_extension",
     "windows_cc_shared_mlir_library",
 )
-load("//jaxlib:pywrap.bzl", "nanobind_pywrap_extension")
 load("//jaxlib:symlink_files.bzl", "symlink_inputs")
 
 package(
@@ -43,7 +44,7 @@ LINKOPTS = select({
     ],
 })
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_mlir",
     srcs = [
         "@llvm-project//mlir:lib/Bindings/Python/MainModule.cpp",
@@ -51,13 +52,14 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "@llvm-project//mlir:MLIRBindingsPythonCore",
+        ":jaxlib_mlir_capi_shared_library",
+        "@llvm-project//mlir:MLIRBindingsPythonCoreNoCAPI",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@nanobind",
     ],
 )
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_mlirDialectsGPU",
     srcs = [
         "@llvm-project//mlir:lib/Bindings/Python/DialectGPU.cpp",
@@ -65,14 +67,15 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "@llvm-project//mlir:CAPIGPU",
-        "@llvm-project//mlir:CAPIIR",
+        ":jaxlib_mlir_capi_shared_library",
+        "@llvm-project//mlir:CAPIGPUHeaders",
+        "@llvm-project//mlir:CAPIIRHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@nanobind",
     ],
 )
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_mlirGPUPasses",
     srcs = [
         "@llvm-project//mlir:lib/Bindings/Python/GPUPasses.cpp",
@@ -80,13 +83,14 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "@llvm-project//mlir:CAPIGPU",
+        ":jaxlib_mlir_capi_shared_library",
+        "@llvm-project//mlir:CAPIGPUHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@nanobind",
     ],
 )
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_mlirDialectsNVGPU",
     srcs = [
         "@llvm-project//mlir:lib/Bindings/Python/DialectNVGPU.cpp",
@@ -94,14 +98,15 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "@llvm-project//mlir:CAPIIR",
-        "@llvm-project//mlir:CAPINVGPU",
+        ":jaxlib_mlir_capi_shared_library",
+        "@llvm-project//mlir:CAPIIRHeaders",
+        "@llvm-project//mlir:CAPINVGPUHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@nanobind",
     ],
 )
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_mlirDialectsLLVM",
     srcs = [
         "@llvm-project//mlir:lib/Bindings/Python/DialectLLVM.cpp",
@@ -109,14 +114,15 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "@llvm-project//mlir:CAPIIR",
-        "@llvm-project//mlir:CAPILLVM",
+        ":jaxlib_mlir_capi_shared_library",
+        "@llvm-project//mlir:CAPIIRHeaders",
+        "@llvm-project//mlir:CAPILLVMHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@nanobind",
     ],
 )
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_mlirDialectsSparseTensor",
     srcs = [
         "@llvm-project//mlir:lib/Bindings/Python/DialectSparseTensor.cpp",
@@ -124,13 +130,14 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "@llvm-project//mlir:CAPISparseTensor",
+        ":jaxlib_mlir_capi_shared_library",
+        "@llvm-project//mlir:CAPISparseTensorHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@nanobind",
     ],
 )
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_mlirSparseTensorPasses",
     srcs = [
         "@llvm-project//mlir:lib/Bindings/Python/SparseTensorPasses.cpp",
@@ -138,20 +145,22 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "@llvm-project//mlir:CAPISparseTensor",
+        ":jaxlib_mlir_capi_shared_library",
+        "@llvm-project//mlir:CAPISparseTensorHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@nanobind",
     ],
 )
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_mosaic_gpu_ext",
     srcs = ["mosaic_gpu_ext.cc"],
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "//jaxlib/mosaic/dialect/gpu:gpu_dialect_capi",
-        "@llvm-project//mlir:CAPIIR",
+        ":jaxlib_mlir_capi_shared_library",
+        "//jaxlib/mosaic/dialect/gpu:gpu_dialect_capi_headers",
+        "@llvm-project//mlir:CAPIIRHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeadersAndDeps",
         "@nanobind",
     ],
@@ -162,16 +171,17 @@ nanobind_pywrap_extension(
 # :jaxlib_mlir_capi_shared_library). This ensures that the RPATH works correctly
 # across platforms. It's not clear if Windows supports RPATH-like functionality
 # across different directories at all.
-nanobind_pywrap_extension(
+py_extension(
     name = "_tpu_ext",
     srcs = ["tpu_ext.cc"],
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "//jaxlib/mosaic:tpu_dialect_capi",
+        ":jaxlib_mlir_capi_shared_library",
+        "//jaxlib/mosaic:tpu_dialect_capi_headers",
         "@com_google_absl//absl/log:check",
         "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:CAPIIR",
+        "@llvm-project//mlir:CAPIIRHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeadersAndDeps",
         "@nanobind",
         "@xla//xla/python:nb_numpy",
@@ -180,7 +190,7 @@ nanobind_pywrap_extension(
 )
 
 # This target contains the extension and it's Python dependencies, which are not
-# supported by the `nanobind_pywrap_extension`/`nanobind_extension` macros.
+# supported by the `py_extension`/`nanobind_extension` macros.
 py_library(
     name = "_tpu_ext_lib",
     deps = [
@@ -190,22 +200,19 @@ py_library(
     ],
 )
 
-nanobind_pywrap_extension(
+nanobind_extension(
     name = "_triton_ext",
     srcs = ["triton_ext.cc"],
     copts = COPTS,
     linkopts = LINKOPTS,
     pytype_srcs = ["_triton_ext.pyi"],
     deps = [
+        ":jaxlib_mlir_capi_shared_library",
+        "//jaxlib/triton:triton_dialect_capi_headers",
+        "@llvm-project//mlir:CAPIIRHeaders",
+        "@llvm-project//mlir:MLIRBindingsPythonNanobindHeadersAndDeps",
         "@nanobind",
-    ] + if_windows(
-        [],
-        [
-            "//jaxlib/triton:triton_dialect_capi",
-            "@llvm-project//mlir:CAPIIR",
-            "@llvm-project//mlir:MLIRBindingsPythonNanobindHeadersAndDeps",
-        ],
-    ),
+    ],
 )
 
 symlink_inputs(
@@ -228,7 +235,7 @@ cc_library(
     hdrs = ["jaxlib_mlir_capi_shims.h"],
     deps = [
         "@llvm-project//mlir:BuiltinToLLVMIRTranslation",
-        "@llvm-project//mlir:CAPIIR",
+        "@llvm-project//mlir:CAPIIRHeaders",
         "@llvm-project//mlir:GPUPipelines",
         "@llvm-project//mlir:GPUToLLVMIRTranslation",
         "@llvm-project//mlir:LLVMToLLVMIRTranslation",
@@ -243,33 +250,34 @@ cc_library(
     name = "jaxlib_mlir_capi_shims_hdrs",
     hdrs = ["jaxlib_mlir_capi_shims.h"],
     deps = [
-        "@llvm-project//mlir:CAPIIR",
+        "@llvm-project//mlir:CAPIIRHeaders",
     ],
 )
 
 # JAX-specific registrations.
-nanobind_pywrap_extension(
+py_extension(
     name = "register_jax_dialects",
     srcs = ["register_jax_dialects.cc"],
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "//jaxlib/mosaic/gpu:mlir_capi",
-        "@llvm-project//mlir:CAPIArith",
-        "@llvm-project//mlir:CAPIGPU",
-        "@llvm-project//mlir:CAPIIR",
-        "@llvm-project//mlir:CAPILLVM",
-        "@llvm-project//mlir:CAPIMath",
-        "@llvm-project//mlir:CAPIMemRef",
-        "@llvm-project//mlir:CAPINVGPU",
-        "@llvm-project//mlir:CAPINVVM",
-        "@llvm-project//mlir:CAPISCF",
-        "@llvm-project//mlir:CAPITransforms",
-        "@llvm-project//mlir:CAPIVector",
+        ":jaxlib_mlir_capi_shared_library",
+        "//jaxlib/mosaic/gpu:mlir_capi_headers",
+        "@llvm-project//mlir:CAPIArithHeaders",
+        "@llvm-project//mlir:CAPIGPUHeaders",
+        "@llvm-project//mlir:CAPIIRHeaders",
+        "@llvm-project//mlir:CAPILLVMHeaders",
+        "@llvm-project//mlir:CAPIMathHeaders",
+        "@llvm-project//mlir:CAPIMemRefHeaders",
+        "@llvm-project//mlir:CAPINVGPUHeaders",
+        "@llvm-project//mlir:CAPINVVMHeaders",
+        "@llvm-project//mlir:CAPISCFHeaders",
+        "@llvm-project//mlir:CAPITransformsHeaders",
+        "@llvm-project//mlir:CAPIVectorHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@local_config_python//:headers",
         "@nanobind",
-        "@shardy//shardy/integrations/c:sdy_capi",
+        "@shardy//shardy/integrations/c:sdy_capi_headers",
     ],
 )
 
@@ -277,7 +285,7 @@ nanobind_pywrap_extension(
 # MHLO Extensions
 ##---------------------------------------------------------------------------##
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_mlirHlo",
     srcs = [
         "@xla//xla/mlir_hlo:bindings/python/MlirHloModule.cc",
@@ -285,11 +293,12 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "@llvm-project//mlir:CAPIIR",
+        ":jaxlib_mlir_capi_shared_library",
+        "@llvm-project//mlir:CAPIIRHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@local_config_python//:headers",
         "@nanobind",
-        "@xla//xla/mlir_hlo:CAPI",
+        "@xla//xla/mlir_hlo:CAPIHeaders",
     ],
 )
 
@@ -297,7 +306,7 @@ nanobind_pywrap_extension(
 # Shardy Extensions
 ##---------------------------------------------------------------------------##
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_sdy",
     srcs = [
         "@shardy//shardy/integrations/python/ir:sdy_module.cc",
@@ -305,12 +314,13 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "@llvm-project//mlir:CAPIIR",
+        ":jaxlib_mlir_capi_shared_library",
+        "@llvm-project//mlir:CAPIIRHeaders",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@local_config_python//:headers",
         "@nanobind",
-        "@shardy//shardy/integrations/c:sdy_capi",
+        "@shardy//shardy/integrations/c:sdy_capi_headers",
     ],
 )
 
@@ -318,7 +328,7 @@ nanobind_pywrap_extension(
 # Stablehlo Extensions
 ##---------------------------------------------------------------------------##
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_chlo",
     srcs = [
         "@stablehlo//:chlo_py_api_files",
@@ -326,15 +336,16 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
-        "@llvm-project//mlir:CAPIIR",
+        ":jaxlib_mlir_capi_shared_library",
+        "@llvm-project//mlir:CAPIIRHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@local_config_python//:headers",
         "@nanobind",
-        "@stablehlo//:chlo_capi",
+        "@stablehlo//:chlo_capi_headers",
     ],
 )
 
-nanobind_pywrap_extension(
+py_extension(
     name = "_stablehlo",
     srcs = [
         "@stablehlo//:stablehlo_py_api_files",
@@ -342,12 +353,13 @@ nanobind_pywrap_extension(
     copts = COPTS,
     linkopts = LINKOPTS,
     deps = [
+        ":jaxlib_mlir_capi_shared_library",
         "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:CAPIIR",
+        "@llvm-project//mlir:CAPIIRHeaders",
         "@llvm-project//mlir:MLIRBindingsPythonNanobindHeaders",
         "@local_config_python//:headers",
         "@nanobind",
-        "@stablehlo//:stablehlo_capi",
+        "@stablehlo//:stablehlo_capi_headers",
     ],
 )
 
diff --git a/jaxlib/mlir/_mlir_libs/triton_ext.cc b/jaxlib/mlir/_mlir_libs/triton_ext.cc
index 687ceec4c..7fba7e1df 100644
--- a/jaxlib/mlir/_mlir_libs/triton_ext.cc
+++ b/jaxlib/mlir/_mlir_libs/triton_ext.cc
@@ -13,8 +13,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef _WIN32
-
 #include <cstdint>
 #include <optional>
 
@@ -76,11 +74,3 @@ NB_MODULE(_triton_ext, m) {
           return encoding;
         });
 }
-
-#else  // _WIN32
-
-#include "nanobind/nanobind.h"
-
-NB_MODULE(_triton_ext, m) {}
-
-#endif  // _WIN32
diff --git a/jaxlib/pyinit_stub.c b/jaxlib/pyinit_stub.c
deleted file mode 100644
index 7fc873d9a..000000000
--- a/jaxlib/pyinit_stub.c
+++ /dev/null
@@ -1,28 +0,0 @@
-/* Copyright 2025 The JAX Authors.
-
-Licensed under the Apache License, Version 2.0 (the "License");
-you may not use this file except in compliance with the License.
-You may obtain a copy of the License at
-
-    http://www.apache.org/licenses/LICENSE-2.0
-
-Unless required by applicable law or agreed to in writing, software
-distributed under the License is distributed on an "AS IS" BASIS,
-WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-See the License for the specific language governing permissions and
-limitations under the License.
-==============================================================================*/
-
-// Stub that reexports Wrapped_PyInit_module as PyInit_module.
-
-extern void* Wrapped_PyInit_@MODULE_NAME@();
-
-#if defined(WIN32) || defined(_WIN32)
-#define EXPORT_SYMBOL __declspec(dllexport)
-#else
-#define EXPORT_SYMBOL __attribute__ ((visibility("default")))
-#endif
-
-EXPORT_SYMBOL void* PyInit_@MODULE_NAME@() {
-  return Wrapped_PyInit_@MODULE_NAME@();
-}
diff --git a/jaxlib/pywrap.bzl b/jaxlib/pywrap.bzl
deleted file mode 100644
index 75324e019..000000000
--- a/jaxlib/pywrap.bzl
+++ /dev/null
@@ -1,89 +0,0 @@
-# Copyright 2025 The JAX Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Wrappers around pywrap rules for JAX."""
-
-load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
-load(
-    "@xla//third_party/py/rules_pywrap:pywrap.impl.bzl",
-    "pybind_extension",
-    _pywrap_binaries = "pywrap_binaries",
-    _pywrap_library = "pywrap_library",
-)
-
-pywrap_library = _pywrap_library
-pywrap_binaries = _pywrap_binaries
-
-def nanobind_pywrap_extension(
-        name,
-        srcs = [],
-        deps = [],
-        pytype_srcs = [],
-        pytype_deps = [],
-        copts = [],
-        linkopts = [],
-        visibility = None):
-    # buildifier: disable=function-docstring-args
-    "Python extension rule using nanobind and the pywrap rules."
-    module_name = name
-    lib_name = name + "_pywrap_library"
-    src_cc_name = name + "_pywrap_stub.c"
-
-    # We put the entire contents of the extension in a single cc_library, which will become part of
-    # the common pywrap library. All the contents of all extensions will end up in the common
-    # library.
-    native.cc_library(
-        name = lib_name,
-        srcs = srcs,
-        copts = copts,
-        deps = deps,
-        local_defines = [
-            "PyInit_{}=Wrapped_PyInit_{}".format(module_name, module_name),
-        ],
-        visibility = ["//visibility:private"],
-    )
-
-    # We build a small stub library as the extension that forwards to the PyInit_... symbol from the
-    # common pywrap library.
-    expand_template(
-        name = name + "_pywrap_stub",
-        testonly = True,
-        out = src_cc_name,
-        substitutions = {
-            "@MODULE_NAME@": module_name,
-        },
-        template = "//jaxlib:pyinit_stub.c",
-        visibility = ["//visibility:private"],
-    )
-
-    # Despite its name "pybind_extension" has nothing to do with pybind. It is the Python extension
-    # rule from the pywrap rules.
-    pybind_extension(
-        name = name,
-        srcs = [src_cc_name],
-        deps = [":" + lib_name],
-        linkopts = linkopts,
-        visibility = visibility,
-        default_deps = [],
-        common_lib_packages = [
-            "jaxlib",
-        ],
-    )
-
-    # Create a py_library with the type stubs as data, on which wheel builds can depend.
-    native.py_library(
-        name = name + "_type_stubs",
-        data = pytype_srcs,
-        deps = pytype_deps,
-    )
diff --git a/jaxlib/setup.py b/jaxlib/setup.py
index 5bd010525..b3a37a25f 100644
--- a/jaxlib/setup.py
+++ b/jaxlib/setup.py
@@ -76,8 +76,6 @@ setup(
     package_data={
         'jaxlib': [
             '*.so',
-            '*.dylib',
-            '*.dll',
             '*.pyd*',
             'py.typed',
             'cpu/*',
diff --git a/jaxlib/tools/BUILD.bazel b/jaxlib/tools/BUILD.bazel
index 3ea09802d..79a1f7e70 100644
--- a/jaxlib/tools/BUILD.bazel
+++ b/jaxlib/tools/BUILD.bazel
@@ -64,10 +64,9 @@ py_binary(
         "LICENSE.txt",
         "//jaxlib",
         "//jaxlib:README.md",
-        "//jaxlib:jaxlib_binaries",
         "//jaxlib:setup.py",
         "//jaxlib/xla:xla_client.py",
-        "//jaxlib/xla:xla_extension_type_stubs",
+        "//jaxlib/xla:xla_extension",
         "@xla//xla/ffi/api:api.h",
         "@xla//xla/ffi/api:c_api.h",
         "@xla//xla/ffi/api:ffi.h",
diff --git a/jaxlib/tools/build_wheel.py b/jaxlib/tools/build_wheel.py
index bab1c6014..fcc811789 100644
--- a/jaxlib/tools/build_wheel.py
+++ b/jaxlib/tools/build_wheel.py
@@ -60,11 +60,11 @@ args = parser.parse_args()
 
 r = runfiles.Create()
 
+
 def _is_mac():
   return platform.system() == "Darwin"
 
 
-soext = "dll" if build_utils.is_windows() else ("dylib" if _is_mac() else "so")
 pyext = "pyd" if build_utils.is_windows() else "so"
 
 
@@ -135,7 +135,7 @@ def verify_mac_libraries_dont_reference_chkstack():
   if not _is_mac():
     return
   nm = subprocess.run(
-      ["nm", "-g", r.Rlocation(f"__main__/jaxlib/xla_extension.{pyext}")],
+      ["nm", "-g", r.Rlocation("__main__/jaxlib/xla/xla_extension.so")],
       capture_output=True,
       text=True,
       check=False,
@@ -186,7 +186,6 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu):
       src_files=[
           f"__main__/jaxlib/cpu_feature_guard.{pyext}",
           f"__main__/jaxlib/utils.{pyext}",
-          "__main__/jaxlib/jax_common.dll" if build_utils.is_windows() else f"__main__/jaxlib/libjax_common.{soext}",
           "__main__/jaxlib/lapack.py",
           "__main__/jaxlib/hlo_helpers.py",
           "__main__/jaxlib/gpu_prng.py",
@@ -199,7 +198,7 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu):
           "__main__/jaxlib/plugin_support.py",
           "__main__/jaxlib/version.py",
           "__main__/jaxlib/xla/xla_client.py",
-          f"__main__/jaxlib/xla_extension.{pyext}",
+          f"__main__/jaxlib/xla/xla_extension.{pyext}",
       ],
   )
   # This file is required by PEP-561. It marks jaxlib as package containing
@@ -312,31 +311,38 @@ def prepare_wheel(sources_path: pathlib.Path, *, cpu):
   )
 
 
+  if build_utils.is_windows():
+    capi_so = "__main__/jaxlib/mlir/_mlir_libs/jaxlib_mlir_capi.dll"
+  else:
+    so_ext = "dylib" if _is_mac() else "so"
+    capi_so = f"__main__/jaxlib/mlir/_mlir_libs/libjaxlib_mlir_capi.{so_ext}"
+
   mlir_libs_dir = jaxlib_dir / "mlir" / "_mlir_libs"
   copy_runfiles(
       dst_dir=mlir_libs_dir,
       src_files=[
+          capi_so,
           "__main__/jaxlib/mlir/_mlir_libs/__init__.py",
-          f"__main__/jaxlib/_mlir.{pyext}",
-          f"__main__/jaxlib/_chlo.{pyext}",
-          f"__main__/jaxlib/_mlirHlo.{pyext}",
-          f"__main__/jaxlib/_mlirDialectsSparseTensor.{pyext}",
-          f"__main__/jaxlib/_mlirSparseTensorPasses.{pyext}",
-          f"__main__/jaxlib/_mosaic_gpu_ext.{pyext}",
-          f"__main__/jaxlib/_tpu_ext.{pyext}",
-          f"__main__/jaxlib/_sdy.{pyext}",
-          f"__main__/jaxlib/_stablehlo.{pyext}",
-          f"__main__/jaxlib/register_jax_dialects.{pyext}",
-          f"__main__/jaxlib/_mlirDialectsGPU.{pyext}",
-          f"__main__/jaxlib/_mlirDialectsLLVM.{pyext}",
-          f"__main__/jaxlib/_mlirDialectsNVGPU.{pyext}",
-          f"__main__/jaxlib/_mlirGPUPasses.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_mlir.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_chlo.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_mlirHlo.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_mlirDialectsSparseTensor.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_mlirSparseTensorPasses.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_mosaic_gpu_ext.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_tpu_ext.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_sdy.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_stablehlo.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/register_jax_dialects.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_mlirDialectsGPU.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_mlirDialectsLLVM.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_mlirDialectsNVGPU.{pyext}",
+          f"__main__/jaxlib/mlir/_mlir_libs/_mlirGPUPasses.{pyext}",
       ]
       + (
           []
           if build_utils.is_windows()
           else [
-              f"__main__/jaxlib/_triton_ext.{pyext}",
+              f"__main__/jaxlib/mlir/_mlir_libs/_triton_ext.{pyext}",
               "__main__/jaxlib/mlir/_mlir_libs/_triton_ext.pyi",
           ]
       ),
diff --git a/jaxlib/triton/BUILD b/jaxlib/triton/BUILD
index 64410fdfe..99cddd9e6 100644
--- a/jaxlib/triton/BUILD
+++ b/jaxlib/triton/BUILD
@@ -35,10 +35,7 @@ pytype_strict_library(
         "//jaxlib/mlir:ir",
     ] + if_windows(
         [],
-        [
-            "//jaxlib/mlir/_mlir_libs:_triton_ext",
-            "//jaxlib/mlir/_mlir_libs:_triton_ext_type_stubs",
-        ],
+        ["//jaxlib/mlir/_mlir_libs:_triton_ext"],
     ),
 )
 
diff --git a/jaxlib/xla/BUILD b/jaxlib/xla/BUILD
index a299629c3..c5d151b2c 100644
--- a/jaxlib/xla/BUILD
+++ b/jaxlib/xla/BUILD
@@ -24,7 +24,6 @@ load(
     "py_strict_test",
     "pytype_strict_library",
 )
-load("//jaxlib:pywrap.bzl", "nanobind_pywrap_extension")
 
 licenses(["notice"])
 
@@ -40,7 +39,7 @@ package_group(
     ],
 )
 
-nanobind_pywrap_extension(
+nanobind_extension(
     name = "xla_extension",
     srcs = ["xla.cc"],
     pytype_deps = py_deps(["numpy"]),


Test patch:
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 20c998d6a..1670f1ee4 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -96,7 +96,7 @@ class NumpyLinalgTest(jtu.JaxTestCase):
       a = rng(factor_shape, dtype)
       return [np.matmul(a, jnp.conj(T(a)))]
 
-    jnp_fun = partial(jnp.linalg.cholesky, upper=upper)
+    jnp_fun = partial(jnp.linalg.cholesky, upper=upper, symmetrize_input=True)
 
     def np_fun(x, upper=upper):
       # Upper argument added in NumPy 2.0.0