code: def _closed_call_discharge_rule(
      in_avals: Sequence[core.AbstractValue],
      _,
      *args,
      call_jaxpr: core.ClosedJaxpr,
  ):
      (
          discharged_closed_jaxpr,
          num_outs,
          fun,
      ) = _cached_closed_jaxpr_discharge(call_jaxpr)
      out_and_ref_vals = core.closed_call_p.bind(
          fun, *args, call_jaxpr=discharged_closed_jaxpr
      )
      out_vals, ref_vals = split_list(out_and_ref_vals, [num_outs])
      ref_vals_iter = iter(ref_vals)
      new_invals = tuple(
          next(ref_vals_iter) if isinstance(aval, AbstractRef) else None
          for aval in in_avals
      )
      sentinel = object()
      assert next(ref_vals_iter, sentinel) is sentinel
      return new_invals, out_vals
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import dataclasses
  from functools import partial
  import operator
  from typing import Any, Callable, Protocol
  import numpy
  from jax._src import api_util
  from jax._src import ad_util
  from jax._src import config
  from jax._src import core
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src import tree_util
  from jax._src.interpreters import ad
  from jax._src.interpreters import mlir
  from jax._src.interpreters import partial_eval
  from jax._src.lax import lax
  from jax._src.lax import slicing
  from jax._src.state import indexing
  from jax._src.state.types import AbstractRef, RefEffect
  from jax._src.state.primitives import get_p, swap_p, addupdate_p
  from jax._src.state.utils import hoist_consts_to_refs
  from jax._src.typing import Array
  from jax._src.util import safe_map, safe_zip, split_list, weakref_lru_cache, partition_list, merge_lists, split_dict
label_desc: Use image node_id 18 for calling a global function with example usage: _closed_call_discharge_rule(in_avals, _) and returns: new_invals,  out_vals