code: def _scan(
      f: Callable[[_Carry, _Input], tuple[_Carry, _Output]],
      init: _Carry,
      xs: Iterable[_Input],
  ) -> tuple[_Carry, _Output]:
      """Implements an unrolled version of scan.
  
      Based on :obj:`jax.lax.scan` and has a similar API.
  
      TODO(schsam): We introduce this function because lax.scan currently has a
      higher peak memory usage than the unrolled version. We will aim to swap this
      out for lax.scan when issue #1273 and related have been resolved.
      """
      carry = init
      ys = []
      flat_xs, tree_def = tree_flatten(xs)
      for flat_x in zip(*flat_xs):
          x = tree_unflatten(tree_def, flat_x)
          carry, y = f(carry, x)
          ys += [y]
  
      return carry, tree_map(lambda *y: jnp.stack(y), *ys)
  
imports: from functools import partial
  from typing import Any, Callable, Iterable, Optional, TypeVar, Union
  import warnings
  import jax
  from jax import device_put
  from jax import devices
  from jax import jit
  from jax import pmap
  from jax import random
  import jax.numpy
  from jax.tree_util import tree_all
  from jax.tree_util import tree_flatten
  from jax.tree_util import tree_map
  from jax.tree_util import tree_unflatten
  import numpy
  from .utils import utils
  from .utils.kernel import Kernel
  from .utils.typing import KernelFn
  from .utils.typing import NTTree
label_desc: Use image node_id 2 for calling a global function with example usage: _scan(f, init, xs) and returns: carry,  tree_map