code: def _for_discharge_rule(
      in_avals,
      _,
      *args: Any,
      jaxpr: core.Jaxpr,
      reverse: bool,
      which_linear: Sequence[bool],
      nsteps: int,
      unroll: int,
  ) -> tuple[Sequence[Any | None], Sequence[Any]]:
      out_vals = for_p.bind(
          *args,
          jaxpr=jaxpr,
          reverse=reverse,
          which_linear=which_linear,
          nsteps=nsteps,
          unroll=unroll,
      )
      new_invals = []
      for aval, out_val in zip(in_avals, out_vals):
          new_invals.append(
              out_val if isinstance(aval, AbstractRef) else None
          )
      return new_invals, out_vals
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  import functools
  import operator
  from typing import Any, Callable, Generic, TypeVar
  import jax.numpy
  from jax import lax
  from jax.api_util import flatten_fun_nokwargs
  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.tree_util import tree_flatten, tree_structure, tree_unflatten, treedef_tuple, tree_map, tree_leaves, PyTreeDef
  from jax._src import ad_util
  from jax._src import core
  from jax._src import dispatch
  from jax._src import dtypes
  from jax._src import linear_util
  from jax._src import source_info_util
  from jax._src.state.types import ReadEffect, AbstractRef, StateEffect
  from jax._src.state import discharge
  from jax._src.state import primitives
  from jax._src.state import utils
  from jax._src.state import types
  from jax._src.typing import Array
  from jax._src.util import partition_list, merge_lists, safe_map, safe_zip, split_list, split_dict, weakref_lru_cache
  from jax._src.lax.control_flow import loops
  from jax._src.lax.control_flow.common import _abstractify, _initial_style_jaxpr
label_desc: Use image node_id 6 for calling a global function with example usage: _for_discharge_rule(in_avals, _) and returns: new_invals,  out_vals