Split weakref_lru_cache into its own extension.

Now that jax-ml@db11efa has landed, we're free to split up xla_extension without creating binary size problems or having to be quite so careful about cross-module dependencies. Here weakref_lru_cache has absolutely nothing to do with XLA.

There's no reason weakref_lru_cache is in the same Python extension as everything else.

PiperOrigin-RevId: 745271825

Author: hawkinsp

jax/_src/lib/BUILD

@@ -44,6 +44,7 @@ py_library_providing_imports_info(
         "//jaxlib/mosaic/python:tpu_dialect",
         "//jaxlib:cpu_feature_guard",
         "//jaxlib:utils",
+        "//jaxlib:weakref_lru_cache",
         "//jaxlib/xla:xla_client",
         "//jaxlib/xla:xla_extension",
         "//jaxlib/triton",

jax/_src/lib/__init__.py

@@ -90,10 +90,23 @@ def _parse_version(v: str) -> tuple[int, ...]:
 from jaxlib.xla_extension import jax_jit as jax_jit  # noqa: F401
 from jaxlib.xla_extension import pmap_lib as pmap_lib  # noqa: F401
 from jaxlib.xla_extension import pytree as pytree  # noqa: F401
-import jaxlib.xla_client as xla_client  # noqa: F401
 
 from jaxlib.xla_extension import Device as Device  # noqa: F401
 
+import jaxlib.xla_client as xla_client  # noqa: F401
+
+# Jaxlib code is split between the Jax and the XLA repositories.
+# Only for the internal usage of the JAX developers, we expose a version
+# number that can be used to perform changes without breaking the main
+# branch on the Jax github.
+jaxlib_extension_version: int = getattr(xla_client, '_version', 0)
+ifrt_version: int = getattr(xla_client, '_ifrt_version', 0)
+
+# TODO(phawkins): remove type: ignore once the minimum jaxlib is bumped.
+if jaxlib_extension_version >= 328:
+  import jaxlib.weakref_lru_cache as weakref_lru_cache  # type: ignore  # noqa: F401
+else:
+  weakref_lru_cache = xla_extension  # type: ignore  # noqa: F401
 
 # XLA garbage collection: see https://github.com/jax-ml/jax/issues/14882
 def _xla_gc_callback(*args):
@@ -113,13 +126,6 @@ def _xla_gc_callback(*args):
 import jaxlib.gpu_prng as gpu_prng  # pytype: disable=import-error  # noqa: F401
 import jaxlib.gpu_linalg as gpu_linalg  # pytype: disable=import-error  # noqa: F401
 
-# Jaxlib code is split between the Jax and the Tensorflow repositories.
-# Only for the internal usage of the JAX developers, we expose a version
-# number that can be used to perform changes without breaking the main
-# branch on the Jax github.
-jaxlib_extension_version: int = getattr(xla_client, '_version', 0)
-ifrt_version: int = getattr(xla_client, '_ifrt_version', 0)
-
 import jaxlib.gpu_rnn as gpu_rnn  # pytype: disable=import-error  # noqa: F401
 import jaxlib.gpu_triton as gpu_triton # pytype: disable=import-error  # noqa: F401
 

jax/_src/util.py

@@ -27,7 +27,7 @@
 import numpy as np
 
 from jax._src import config
-from jax._src.lib import xla_client as xc
+from jax._src.lib import weakref_lru_cache as _weakref_lru_cache
 from jax._src.lib import utils as jaxlib_utils
 
 logger = logging.getLogger(__name__)
@@ -304,8 +304,9 @@ def weakref_lru_cache(call: Callable, maxsize=2048,
   behave similar to `functools.lru_cache`.
   """
   global _weakref_lru_caches
-  cached_call = xc.weakref_lru_cache(
-      config.trace_context if trace_context_in_key else _ignore, call, maxsize)
+  cached_call = _weakref_lru_cache.weakref_lru_cache(
+      config.trace_context if trace_context_in_key else _ignore, call, maxsize
+  )
   _weakref_lru_caches.add(cached_call)
   return cached_call
 

jaxlib/BUILD

@@ -17,7 +17,9 @@
 load(
     "//jaxlib:jax.bzl",
     "nanobind_extension",
+    "py_deps",
     "py_library_providing_imports_info",
+    "py_strict_test",
     "pytype_library",
 )
 load(
@@ -59,6 +61,7 @@ py_library_providing_imports_info(
         ":cpu_feature_guard",
         ":jax",
         ":utils",
+        ":weakref_lru_cache",
         "//jaxlib/cpu:_lapack",
         "//jaxlib/mlir",
         "//jaxlib/mlir:arithmetic_dialect",
@@ -119,6 +122,7 @@ pywrap_library(
     },
     deps = [
         ":utils",
+        ":weakref_lru_cache",
         "//jaxlib/mlir/_mlir_libs:_chlo",
         "//jaxlib/mlir/_mlir_libs:_mlir",
         "//jaxlib/mlir/_mlir_libs:_mlirDialectsGPU",
@@ -215,6 +219,35 @@ nanobind_extension(
     ],
 )
 
+nanobind_pywrap_extension(
+    name = "weakref_lru_cache",
+    srcs = ["weakref_lru_cache.cc"],
+    pytype_srcs = ["weakref_lru_cache.pyi"],
+    deps = [
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/cleanup",
+        "@com_google_absl//absl/hash",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/synchronization",
+        "@nanobind",
+        "@xla//third_party/python_runtime:headers",
+        "@xla//xla/pjrt:lru_cache",
+        "@xla//xla/tsl/platform:logging",
+    ],
+)
+
+py_strict_test(
+    name = "weakref_lru_cache_test",
+    srcs = ["weakref_lru_cache_test.py"],
+    deps = [
+        ":weakref_lru_cache",
+    ] + py_deps([
+        "absl/flags",
+        "absl/logging",
+        "absl/testing",
+    ]),
+)
+
 nanobind_pywrap_extension(
     name = "utils",
     srcs = ["utils.cc"],

jaxlib/tools/build_wheel.py

@@ -103,7 +103,6 @@ def patch_copy_mlir_import(src_file, dst_dir):
     "pytree.pyi",
     "transfer_guard_lib.pyi",
 ]
-_OPTIONAL_XLA_EXTENSION_STUBS = []
 
 
 def patch_copy_xla_extension_stubs(dst_dir):
@@ -112,8 +111,6 @@ def patch_copy_xla_extension_stubs(dst_dir):
   for stub_name in _XLA_EXTENSION_STUBS:
     stub_path = r.Rlocation("__main__/jaxlib/xla/xla_extension/" + stub_name)
     stub_path = str(stub_path)  # Make pytype accept os.path.exists(stub_path).
-    if stub_name in _OPTIONAL_XLA_EXTENSION_STUBS and not os.path.exists(stub_path):
-      continue
     with open(stub_path) as f:
       src = f.read()
     src = src.replace(
@@ -199,6 +196,8 @@ 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/weakref_lru_cache.{pyext}",
+          "__main__/jaxlib/weakref_lru_cache.pyi",
           f"__main__/jaxlib/xla_extension.{pyext}",
       ],
   )

jaxlib/xla/weakref_lru_cache.cc jaxlib/weakref_lru_cache.ccjaxlib/xla/weakref_lru_cache.cc renamed to jaxlib/weakref_lru_cache.cc

@@ -13,8 +13,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "jaxlib/xla/weakref_lru_cache.h"
-
 #include <Python.h>
 
 #include <cstddef>
@@ -369,7 +367,7 @@ class WeakrefLRUCache : public std::enable_shared_from_this<WeakrefLRUCache> {
     {0, nullptr},
 };
 
-void BuildWeakrefLRUCacheAPI(nb::module_& m) {
+NB_MODULE(weakref_lru_cache, m) {
   auto weakref_lru_cache =
       nb::class_<WeakrefLRUCache>(m, "WeakrefLRUCache",
                                   nb::is_weak_referenceable(),

jaxlib/weakref_lru_cache.pyi

@@ -0,0 +1,37 @@
+# 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.
+# ==============================================================================
+
+from collections.abc import Callable
+
+class WeakrefLRUCache:
+  def __call__(self, arg0: object, /, *args, **kwargs) -> object: ...
+  def cache_keys(self) -> list[object]: ...
+  def cache_info(self) -> WeakrefLRUCache.WeakrefLRUCacheInfo: ...
+  def cache_clear(self) -> None: ...
+
+  class WeakrefLRUCacheInfo:
+    @property
+    def hits(self) -> int: ...
+    @property
+    def misses(self) -> int: ...
+    @property
+    def maxsize(self) -> int: ...
+    @property
+    def currsize(self) -> int: ...
+    def __repr__(self) -> str: ...
+
+def weakref_lru_cache(
+    cache_context_fn: Callable, fn: Callable, maxsize: int = 2048
+) -> WeakrefLRUCache: ...

jaxlib/xla/weakref_lru_cache_test.py jaxlib/weakref_lru_cache_test.pyjaxlib/xla/weakref_lru_cache_test.py renamed to jaxlib/weakref_lru_cache_test.py

@@ -19,8 +19,7 @@
 import weakref
 
 from absl.testing import absltest
-
-from jax.jaxlib.xla import xla_client
+from jax.jaxlib import weakref_lru_cache
 
 
 class WeakrefLRUCacheTest(absltest.TestCase):
@@ -60,7 +59,7 @@ def CacheFn(obj, gil_releasing_cache_key):
       del gil_releasing_cache_key
       return None
 
-    cache = xla_client.weakref_lru_cache(lambda: None, CacheFn, 2048)
+    cache = weakref_lru_cache.weakref_lru_cache(lambda: None, CacheFn, 2048)
 
     wrkey = WRKey()
 
@@ -79,7 +78,9 @@ def Body():
   def testAnotherMultiThreaded(self):
     num_workers = 5
     barrier = threading.Barrier(num_workers)
-    cache = xla_client.weakref_lru_cache(lambda: None, lambda x, y: y, 2048)
+    cache = weakref_lru_cache.weakref_lru_cache(
+        lambda: None, lambda x, y: y, 2048
+    )
 
     class WRKey:
       pass
@@ -118,7 +119,7 @@ def CacheFn(obj, kwkey1, kwkey2):
       miss_id += 1
       return miss_id
 
-    cache = xla_client.weakref_lru_cache(lambda: None, CacheFn, 4)
+    cache = weakref_lru_cache.weakref_lru_cache(lambda: None, CacheFn, 4)
 
     wrkey = WRKey()
 
@@ -131,7 +132,7 @@ def CacheFn(obj, arg):
       del obj
       return arg + "extra"
 
-    cache = xla_client.weakref_lru_cache(lambda: None, CacheFn, 4)
+    cache = weakref_lru_cache.weakref_lru_cache(lambda: None, CacheFn, 4)
 
     class WRKey:
       pass
@@ -151,7 +152,7 @@ class NonWRKey:
     with self.assertRaises(TypeError):
       weakref.ref(non_wr_key)
 
-    cache = xla_client.weakref_lru_cache(lambda: None, lambda x: 2048)
+    cache = weakref_lru_cache.weakref_lru_cache(lambda: None, lambda x: 2048)
     for _ in range(100):
       with self.assertRaises(TypeError):
         cache(non_wr_key)
@@ -169,7 +170,9 @@ def __eq__(self, other):
       def __hash__(self):
         raise ValueError("hash")
 
-    cache = xla_client.weakref_lru_cache(lambda: None, lambda x, y: y, 2048)
+    cache = weakref_lru_cache.weakref_lru_cache(
+        lambda: None, lambda x, y: y, 2048
+    )
     wrkey = WRKey()
     with self.assertRaises(ValueError):
       for _ in range(100):
@@ -179,7 +182,9 @@ def testPrintingStats(self):
     class WRKey:
       pass
 
-    cache = xla_client.weakref_lru_cache(lambda: None, lambda x, y: y, 2048)
+    cache = weakref_lru_cache.weakref_lru_cache(
+        lambda: None, lambda x, y: y, 2048
+    )
     wrkey = WRKey()
     for i in range(10):
       cache(wrkey, i)
@@ -203,7 +208,9 @@ def __eq__(self, other):
       def __hash__(self):
         return hash(self.x)
 
-    cache = xla_client.weakref_lru_cache(lambda: None, lambda x, y: y, 2048)
+    cache = weakref_lru_cache.weakref_lru_cache(
+        lambda: None, lambda x, y: y, 2048
+    )
     keys = [WRKey(i) for i in range(10)]
     for i in range(10):
       cache(keys[i], i)
@@ -225,7 +232,7 @@ def CallFn(x, y, *args, **kwargs):
       del x, args, kwargs
       return y
 
-    cache = xla_client.weakref_lru_cache(CacheContextFn, CallFn, 2048)
+    cache = weakref_lru_cache.weakref_lru_cache(CacheContextFn, CallFn, 2048)
 
     keys = [WRKey() for _ in range(10)]
     values = [str(i) for i in range(10)]
@@ -239,7 +246,7 @@ def CallFn(x, y, *args, **kwargs):
         [
             CacheContextFn,
             CallFn,
-            xla_client._xla.WeakrefLRUCache,
+            weakref_lru_cache.WeakrefLRUCache,
             kwargs,
         ]
         + [weakref.getweakrefs(key)[0] for key in keys]

jaxlib/xla/BUILD

@@ -63,7 +63,6 @@ nanobind_pywrap_extension(
         ":sdy",
         ":traceback",
         ":util",
-        ":weakref_lru_cache",
         ":xla_compiler",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -876,29 +875,6 @@ cc_library(
     ],
 )
 
-cc_library(
-    name = "weakref_lru_cache",
-    srcs = ["weakref_lru_cache.cc"],
-    hdrs = ["weakref_lru_cache.h"],
-    compatible_with = [],
-    copts = [
-        "-fexceptions",
-        "-fno-strict-aliasing",
-    ],
-    features = ["-use_header_modules"],
-    deps = [
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/cleanup",
-        "@com_google_absl//absl/hash",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/synchronization",
-        "@nanobind",
-        "@xla//third_party/python_runtime:headers",
-        "@xla//xla/pjrt:lru_cache",
-        "@xla//xla/tsl/platform:logging",
-    ],
-)
-
 cc_library(
     name = "xla_compiler",
     srcs = ["xla_compiler.cc"],
@@ -1041,18 +1017,6 @@ py_strict_test(
     ]),
 )
 
-py_strict_test(
-    name = "weakref_lru_cache_test",
-    srcs = ["weakref_lru_cache_test.py"],
-    deps = [
-        ":xla_client",
-    ] + py_deps([
-        "absl/flags",
-        "absl/logging",
-        "absl/testing",
-    ]),
-)
-
 py_strict_test(
     name = "pytree_test",
     srcs = ["pytree_test.py"],