Changed how ifrt_proxy is imported to improve backwards compatibility with JAX

lukebaumann · copybara-github · commit 359776d45494 · 2024-09-25T13:25:22.000-07:00
PiperOrigin-RevId: 678824312
diff --git a/pathwaysutils/proxy_backend.py b/pathwaysutils/proxy_backend.py
@@ -15,7 +15,7 @@
 
 import jax
 from jax._src import xla_bridge
-from jax.lib.xla_extension import ifrt_proxy
+from jaxlib.xla_extension import ifrt_proxy
 
 
 def register_backend_factory():
