code: def _cached_closed_jaxpr_discharge(closed_jaxpr):
      jaxpr, consts = closed_jaxpr.jaxpr, closed_jaxpr.consts
      num_outs = len(jaxpr.outvars)
      discharged_jaxpr, discharged_consts = discharge_state(
          jaxpr, consts
      )
      discharged_closed_jaxpr = core.ClosedJaxpr(
          discharged_jaxpr, discharged_consts
      )
      fun = lu.wrap_init(core.jaxpr_as_fun(discharged_closed_jaxpr))
      return discharged_closed_jaxpr, num_outs, fun
  
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 17 for calling a global function with example usage: _cached_closed_jaxpr_discharge(closed_jaxpr) and returns: discharged_closed_jaxpr,  num_outs,  fun