code: def fori_loop(lower, upper, body_fun, init_val):
      """Equivalent to :func:`jax.lax.fori_loop` with Haiku state passed in/out."""
      if not base.inside_transform():
          raise ValueError(
              "hk.fori_loop() should not be used outside of hk.transform(). "
              "Use jax.lax.fori_loop() instead."
          )
  
      @functools.wraps(body_fun)
      def pure_body_fun(i, val):
          state, val = val
          with temporary_internal_state(
              state
          ), base.push_jax_trace_level():
              val = body_fun(i, val)
              reserve_up_to_full_rng_block()
              state = internal_state()
              return state, val
  
      if not base.params_frozen():
          # During init we need to unwind one step of the loop to ensure the Haiku
          # state before and after the body has the same structure.
          init_val = body_fun(lower, init_val)
          lower += 1
          try:
              if upper - lower == 0:
                  return init_val
          except jax.errors.ConcretizationTypeError:
              # upper or lower might be tracers, which jax.lax.fori_loop can handle.
              pass
  
      reserve_up_to_full_rng_block()
      state = internal_state()
      init_val = state, init_val
      state, val = jax.lax.fori_loop(
          lower, upper, pure_body_fun, init_val
      )
      update_internal_state(state)
      return val
  
imports: import collections
  import collections.abc
  from collections.abc import Mapping, MutableMapping
  import functools
  import inspect
  from typing import Any, Callable, Optional, TypeVar
  from haiku._src import base
  import jax
  import jax.numpy
label_desc: Use image node_id 21 for calling a global function with example usage: fori_loop(lower, upper, body_fun, init_val) and returns: val, state,  val, init_val