Fix Python path for executing "hi" JAX code.

PiperOrigin-RevId: 829303995

Author: chr1sj0nes

jax/_src/stages.py

@@ -435,8 +435,14 @@ def fall(self):
     _, closed_over_himutables = pe.convert_const_himutables(hi_jaxpr)
     if closed_over_himutables: raise NotImplementedError  # TODO(mattjj)
     lo_jaxpr = pe.lower_jaxpr(hi_jaxpr)
-    in_tree = lojax_pytree(hi_jaxpr.in_aval_qdds, self._in_tree)
-    out_tree = lojax_pytree(hi_jaxpr.out_avals, self.out_tree)
+    if any(a.is_high for a in hi_jaxpr.final_aval_qdds):
+      in_tree = lojax_pytree(hi_jaxpr.in_aval_qdds, self._in_tree)
+    else:
+      in_tree = self._in_tree
+    if any(a.is_high for a in hi_jaxpr.out_avals):
+      out_tree = lojax_pytree(hi_jaxpr.out_avals, self.out_tree)
+    else:
+      out_tree = self.out_tree
     params = dict(lojax_expand_params(hi_jaxpr, self._params), jaxpr=lo_jaxpr)
     lo_meta_tys = [mty.replace(aval=lo_ty)
                    for mty, aq in zip(self._meta_tys_flat, hi_jaxpr.in_aval_qdds)
@@ -791,7 +797,7 @@ def call(*args, **kwargs):
           "function was compiled by a transformation that does not support "
           f"keyword arguments, but called with keyword arguments: {kws}")
 
-    if params.is_high:
+    if params.in_types and any(a.is_high for a in params.in_types[1]):
       hi_args_flat, in_hi_tree = tree_util.tree_flatten((args, kwargs))
       in_hi_tree_, final_qdds = params.in_types
       args_flat = [a.read_loval(core.cur_qdd(x), x) if (a := typeof(x)).has_qdd
@@ -828,7 +834,7 @@ def call(*args, **kwargs):
               f"Tracer type {type(arg)}.")
     lo_outs = params.executable.call(*params.const_args, *args_flat)
 
-    if params.is_high:
+    if params.out_types and any(a.is_high for a in params.out_types[1]):
       out_mut, lo_outs = util.split_list(lo_outs, [_num_himuts_out(final_qdds)])
       _apply_himut(final_qdds, hi_args_flat, out_mut)
       out_hi_tree, out_hi_types = params.out_types

tests/hijax_test.py

@@ -534,6 +534,7 @@ def f(x):
 
     if jit:
       f = jax.jit(f)
+      self.assertEqual(f.lower(2.0).compile()(2.0), 8.0)
 
     self.assertEqual(f(2.0), 8.0)
     xs = jnp.arange(3.0)