code: def discharged_for_loop(
      nsteps, body, init_state, *, reverse: bool = False
  ):
      """A `for_loop` implementation that discharges its body right away.
  
      Potentially useful for testing and benchmarking.
      """
      flat_state, state_tree = tree_flatten(init_state)
      state_avals = map(state_utils.val_to_ref_aval, flat_state)
      idx_aval = core.ShapedArray((), jnp.dtype("int32"))
      jaxpr, consts, out_tree = _trace_to_jaxpr_with_refs(
          body, state_tree, [idx_aval, *state_avals]
      )
      if out_tree != tree_structure(None):
          raise Exception("`body` should not return anything.")
      discharged_jaxpr, discharged_consts = discharge_state(
          jaxpr, consts
      )
  
      def fori_body(i, carry):
          i = jnp.int32(i)
          if reverse:
              i = nsteps - i - 1
          out_flat = core.eval_jaxpr(
              discharged_jaxpr, discharged_consts, i, *carry
          )
          return out_flat
  
      out_flat = loops.fori_loop(0, nsteps, fori_body, flat_state)
      return tree_unflatten(state_tree, out_flat)
  
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 21 for calling a global function with example usage: discharged_for_loop(nsteps, body, init_state) and returns: tree_unflatten, out_flat