code: def scan(
      f: Callable[[Carry, X], tuple[Carry, Y]],
      init: Carry,
      xs: X,
      length: int | None = None,
      reverse: bool = False,
      unroll: int = 1,
  ) -> tuple[Carry, Y]:
      if not callable(f):
          raise TypeError("scan: f argument should be a callable.")
      if unroll < 1:
          raise ValueError("`unroll` must be a positive integer.")
      xs_flat, xs_tree = tree_flatten(xs)
  
      try:
          lengths = [x.shape[0] for x in xs_flat]
      except AttributeError as err:
          msg = "scan got value with no leading axis to scan over: {}."
          raise ValueError(
              msg.format(
                  ", ".join(
                      str(x) for x in xs_flat if not hasattr(x, "shape")
                  )
              )
          ) from err
  
      if length is not None:
          length = int(length)
          if not all(length == l for l in lengths):
              msg = (
                  "scan got `length` argument of {} which disagrees with "
                  "leading axis sizes {}."
              )
              raise ValueError(
                  msg.format(length, [x.shape[0] for x in xs_flat])
              )
      else:
          unique_lengths = set(lengths)
          if len(unique_lengths) > 1:
              msg = "scan got values with different leading axis sizes: {}."
              raise ValueError(
                  msg.format(
                      ", ".join(str(x.shape[0]) for x in xs_flat)
                  )
              )
          elif len(unique_lengths) == 0:
              msg = "scan got no values to scan over and `length` not provided."
              raise ValueError(msg)
          else:
              (length,) = unique_lengths
  
      x_shapes = [x.shape[1:] for x in xs_flat]
      x_dtypes = [dtypes.canonicalize_dtype(x.dtype) for x in xs_flat]
      x_avals = tuple(map(core.ShapedArray, x_shapes, x_dtypes))
  
      def _create_jaxpr(init):
          init_flat = tree_leaves(init)
          _, in_tree = tree_flatten((init, xs))
  
          carry_avals = tuple(map(_abstractify, init_flat))
          jaxpr, _, out_tree = _initial_style_jaxpr(
              f, in_tree, carry_avals + x_avals, "scan"
          )
          return jaxpr, out_tree
  
      jaxpr, out_tree = _create_jaxpr(init)
      _, ys_avals = tree_unflatten(out_tree, jaxpr.out_avals)
      ys = tree_map(
          lambda aval: jnp.zeros([length, *aval.shape], aval.dtype),
          ys_avals,
      )
  
      def for_body(i, refs):
          carry_refs, xs_refs, ys_refs = refs
          carry = tree_map(lambda x: x[()], carry_refs)
          x = tree_map(lambda x: x[i], xs_refs)
          carry, y = f(carry, x)
          tree_map(
              lambda c_ref, c: ref_set(c_ref, (), c), carry_refs, carry
          )
          tree_map(lambda y_ref, y: ref_set(y_ref, (i,), y), ys_refs, y)
  
      assert isinstance(length, int)
      init, _, ys = for_loop(
          length,
          for_body,
          (init, xs, ys),
          reverse=reverse,
          unroll=unroll,
      )
      return init, ys
  
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 4 for calling a global function with example usage: scan(f, init, xs, length, reverse, unroll) and returns: init,  ys, jaxpr,  out_tree