code: def scan(f, init, xs, length=None, reverse=False, unroll=1):
      """Equivalent to :func:`jax.lax.scan` but with Haiku state passed in/out."""
      if not base.inside_transform():
          raise ValueError(
              "hk.scan() should not be used outside of hk.transform(). "
              "Use jax.lax.scan() instead."
          )
  
      if length is None:
          length = jax.tree_util.tree_leaves(xs)[0].shape[0]
  
      running_init_fn = not base.params_frozen()
  
      if running_init_fn:
          # During `init` we need to unroll one step of the scan, this is because our
          # carry contains the Haiku state and during `init` this may change structure
          # (e.g. as state is created).
          if not length:
              x0 = jax.tree_util.tree_map(
                  lambda x: jnp.zeros(x.shape[1:], x.dtype), xs
              )
              _, y0 = f(init, x0)
              y0 = jax.tree_util.tree_map(
                  lambda y: jnp.zeros((0,) + y.shape, y.dtype), y0
              )
              return init, y0
  
          if reverse:
              x0 = jax.tree_util.tree_map(lambda x: x[-1], xs)
              xs = jax.tree_util.tree_map(lambda x: x[:-1], xs)
          else:
              x0 = jax.tree_util.tree_map(lambda x: x[0], xs)
              xs = jax.tree_util.tree_map(lambda x: x[1:], xs)
          init, y0 = f(init, x0)
          y0 = jax.tree_util.tree_map(
              lambda y: jnp.expand_dims(y, 0), y0
          )
          length -= 1
          if not length:
              return init, y0
  
      @functools.wraps(f)
      def stateful_fun(carry, x):
          carry, state = carry
          with temporary_internal_state(state):
              with base.assert_no_new_parameters(), base.push_jax_trace_level():
                  carry, out = f(carry, x)
              reserve_up_to_full_rng_block()
              carry = (carry, internal_state(params=False))
              return carry, out
  
      # Before pulling out the  internal state,  reserve a full block  of RNG keys.
      # This is to make sure we're always passing in the same amount of subkeys in
      # and out of the scan carry (scan requires equal length lists).
      # After every scan iteration we reserve back up to the full block.
      reserve_up_to_full_rng_block()
  
      # We know that we don't need to thread params in and out, since for init we
      # have already created them (given that above we unroll one step of the scan)
      # and for apply we know they are immutable. As such we only need to thread the
      # state and rng in and out.
  
      init = (init, internal_state(params=False))
      (carry, state), ys = jax.lax.scan(
          stateful_fun, init, xs, length, reverse, unroll=unroll
      )
      update_internal_state(state)
  
      if running_init_fn:
          if reverse:
              ys = jax.tree_util.tree_map(
                  lambda y0, ys: jnp.concatenate([ys, y0]), y0, ys
              )
          else:
              ys = jax.tree_util.tree_map(
                  lambda y0, ys: jnp.concatenate([y0, ys]), y0, ys
              )
  
      return carry, ys
  
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 19 for calling a global function with example usage: scan(f, init, xs, length, reverse, unroll) and returns: carry,  ys, init,  y0, init,  y0, carry,  out