Avoid using relatively new jax.extend package until all modules that we use are part of a stable JAX release.

PiperOrigin-RevId: 676890866

Author: lukebaumann

pathwaysutils/plugin_executable.py

@@ -19,7 +19,7 @@
 
 import jax
 from jax._src.interpreters import pxla
-from jax.extend.ifrt_programs import ifrt_programs
+from jax._src.lib.xla_extension import ifrt_programs
 
 
 class PluginExecutable:

pathwaysutils/proxy_backend.py

@@ -14,12 +14,12 @@
 """Register the IFRT Proxy as a backend for JAX."""
 
 import jax
-from jax.extend import backend
+from jax._src import xla_bridge
 from jax.lib.xla_extension import ifrt_proxy
 
 
 def register_backend_factory():
-  backend.register_backend_factory(
+  xla_bridge.register_backend_factory(
       "proxy",
       lambda: ifrt_proxy.get_client(
           jax.config.read("jax_backend_target"),