Use the internal pxla API explicitly because the public pxla is deprecated as of JAX 0.8.2.

Also updated `jax.extend.ifrt_programs` import to import the module and make an alias for the `jax.extend.ifrt_programs.ifrt_programs` variable.

PiperOrigin-RevId: 852516195

Author: lukebaumann

pathwaysutils/plugin_executable.py

@@ -18,8 +18,11 @@
 from typing import List, Sequence, Tuple, Union
 
 import jax
-from jax.interpreters import pxla
-from jax.extend.ifrt_programs import ifrt_programs
+from jax._src.interpreters import pxla  # pylint: disable=protected-access
+from jax.extend import ifrt_programs
+
+
+ifrt_programs = ifrt_programs.ifrt_programs
 
 
 class PluginExecutable: