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: 857248364

Author: lukebaumann

pathwaysutils/plugin_executable.py

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