Change from using jaxlib.xla_client.ifrt_programs to jax.extend.ifrt_programs.ifrt_programs.

PiperOrigin-RevId: 676551752

Author: lukebaumann

pathwaysutils/plugin_executable.py

@@ -19,16 +19,16 @@
 
 import jax
 from jax._src.interpreters import pxla
-from jaxlib import xla_client
+from jax.extend.ifrt_programs import ifrt_programs
 
 
 class PluginExecutable:
   """Class for running compiled IFRT program over the IFRT Proxy."""
 
   def __init__(self, prog_str: str):
     ifrt_client = jax.local_devices()[0].client
-    program = xla_client.ifrt_programs.make_plugin_program(prog_str)
-    options = xla_client.ifrt_programs.make_plugin_compile_options()
+    program = ifrt_programs.make_plugin_program(prog_str)
+    options = ifrt_programs.make_plugin_compile_options()
     self.compiled = ifrt_client.compile_ifrt_program(program, options)
 
   def call(