Introduce jax shim to support multiple versions of JAX

PiperOrigin-RevId: 789952304

Author: lukebaumann

pathwaysutils/jax/__init__.py

@@ -0,0 +1,49 @@
+# Copyright 2025 Google LLC
+#
+# 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.
+"""Pathways JAX abstractions.
+
+This introduces an abstrction layer some JAX APIs that have changed over
+`pathwaysutils`'s compatibility window.
+"""
+
+from typing import Any
+
+try:
+  # jax>=0.7.0
+  from jax.extend import backend  # pylint: disable=g-import-not-at-top
+
+  register_backend_cache = backend.register_backend_cache
+
+  del backend
+except AttributeError:
+  # jax<0.7.0
+  from jax._src import util  # pylint: disable=g-import-not-at-top
+
+  def register_backend_cache(cache: Any, name: str):  # pylint: disable=unused-argument
+    return util.cache_clearing_funs.add(cache.cache_clear)
+
+  del util
+
+try:
+  # jax>0.7.0
+  from jax.extend import backend  # pylint: disable=g-import-not-at-top
+
+  ifrt_proxy = backend.ifrt_proxy
+  del backend
+except AttributeError:
+  # jax<=0.7.0
+  from jax.lib import xla_extension  # pylint: disable=g-import-not-at-top
+
+  ifrt_proxy = xla_extension.ifrt_proxy
+  del xla_extension

pathwaysutils/lru_cache.py

@@ -16,7 +16,7 @@
 import functools
 from typing import Any, Callable
 
-from jax._src import util
+from pathwaysutils import jax as pw_jax
 
 
 def lru_cache(
@@ -38,7 +38,7 @@ def wrap(f):
 
     wrapper.cache_clear = cached.cache_clear
     wrapper.cache_info = cached.cache_info
-    util.cache_clearing_funs.add(wrapper.cache_clear)
+    pw_jax.register_backend_cache(wrapper, "Pathways LRU cache")
     return wrapper
 
   return wrap

pathwaysutils/proxy_backend.py

@@ -15,15 +15,15 @@
 
 import jax
 from jax.extend import backend
-from jax.lib.xla_extension import ifrt_proxy
+from pathwaysutils import jax as pw_jax
 
 
 def register_backend_factory():
   backend.register_backend_factory(
       "proxy",
-      lambda: ifrt_proxy.get_client(
+      lambda: pw_jax.ifrt_proxy.get_client(
           jax.config.read("jax_backend_target"),
-          ifrt_proxy.ClientConnectionOptions(),
+          pw_jax.ifrt_proxy.ClientConnectionOptions(),
       ),
       priority=-1,
   )

pathwaysutils/test/proxy_backend_test.py

@@ -17,10 +17,9 @@
 
 import jax
 from jax.extend import backend
-from jax.lib.xla_extension import ifrt_proxy
+from pathwaysutils import jax as pw_jax
 from pathwaysutils import proxy_backend
 
-
 from absl.testing import absltest
 
 
@@ -39,7 +38,7 @@ def test_no_proxy_backend_registration_raises_error(self):
   def test_proxy_backend_registration(self):
     self.enter_context(
         mock.patch.object(
-            ifrt_proxy,
+            pw_jax.ifrt_proxy,
             "get_client",
             return_value=mock.MagicMock(),
         )