code: def discharge_state(
      jaxpr: core.Jaxpr,
      consts: Sequence[Any],
      *,
      should_discharge: bool | Sequence[bool] = True,
  ) -> tuple[core.Jaxpr, list[Any]]:
      """Converts a jaxpr that takes in `Ref`s into one that doesn't."""
      if isinstance(should_discharge, bool):
          should_discharge = [should_discharge] * len(jaxpr.invars)
      in_avals = [
          v.aval.inner_aval
          if isinstance(v.aval, AbstractRef) and d
          else v.aval
          for v, d in zip(jaxpr.invars, should_discharge)
      ]
      eval_jaxpr = lu.wrap_init(
          partial(
              _eval_jaxpr_discharge_state,
              jaxpr,
              should_discharge,
              consts,
          )
      )
      new_jaxpr, _, new_consts = pe.trace_to_jaxpr_dynamic(
          eval_jaxpr, in_avals
      )
      return new_jaxpr, new_consts
  
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 1 for calling a global function with example usage: discharge_state(jaxpr, consts) and returns: new_jaxpr,  new_consts