code: def _join_cond_pe_staged_jaxpr_inputs(
      jaxprs, all_res_avals, res_aval_indices_per_jaxpr
  ):
      newvar = core.gensym([j.jaxpr for j in jaxprs], suffix="_")
      all_res_vars = map(newvar, all_res_avals)
  
      def augment_jaxpr(jaxpr, res_indices):
          num_res = len(res_indices)
          res_vars = jaxpr.jaxpr.invars[:num_res]
          non_res_vars = jaxpr.jaxpr.invars[num_res:]
  
          aug_res_vars = list(
              util.subvals(all_res_vars, zip(res_indices, res_vars))
          )
          aug_invars = aug_res_vars + non_res_vars
          jaxpr_aug = core.Jaxpr(
              jaxpr.jaxpr.constvars,
              aug_invars,
              jaxpr.jaxpr.outvars,
              jaxpr.jaxpr.eqns,
              jaxpr.jaxpr.effects,
          )
          jaxpr_aug = core.ClosedJaxpr(jaxpr_aug, jaxpr.consts)
          return jaxpr_aug
  
      return tuple(
          map(augment_jaxpr, jaxprs, res_aval_indices_per_jaxpr)
      )
  
imports: from __future__ import annotations
  import collections
  from collections.abc import Sequence
  import functools
  from functools import partial
  import inspect
  import itertools
  import operator
  from typing import Any, Callable, TypeVar
  from jax.tree_util import tree_flatten, tree_unflatten
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import effects
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import util
  from jax._src.state.discharge import register_discharge_rule, discharge_state
  from jax._src.state.types import AbstractRef, RefEffect
  from jax._src.core import ConcreteArray, raise_to_shaped, replace_jaxpr_effects
  from jax._src.interpreters import ad
  from jax._src.interpreters import batching
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.interpreters import xla
  from jax._src.lax import lax
  from jax._src.traceback_util import api_boundary
  from jax._src.util import safe_map, split_list, partition_list
  from jax._src.lib.mlir import ir
  from jax._src.lib.mlir.dialects import hlo
  import numpy
  from jax._src.lax.control_flow.common import _abstractify, _avals_short, _check_tree_and_avals, _initial_style_jaxprs_with_common_consts, _make_closed_jaxpr, _prune_zeros, _typecheck_param
label_desc: Use image node_id 15 for calling a global function with example usage: _join_cond_pe_staged_jaxpr_inputs(jaxprs, all_res_avals, res_aval_indices_per_jaxpr) and returns: tuple, jaxpr_aug