code: def scan(
      f: Callable[[_Carry, _X], tuple[_Carry, _Y]],
      init: _Carry,
      xs: _X,
      length: Optional[int] = None,
      *,
      buffers: Optional[
          Callable[[_Carry], Union[_Node, Sequence[_Node]]]
      ] = None,
      kind: Literal["lax", "checkpointed"],
      checkpoints: Union[None, int, Literal["all"]] = None,
  ) -> tuple[_Carry, _Y]:
      """As `jax.lax.scan`, but with optional checkpointing to reduce memory usage.
  
      **Arguments:**
  
      - `f`: As `jax.lax.scan`.
      - `init`: As `jax.lax.scan`.
      - `xs`: As `jax.lax.scan`.
      - `length`: As `jax.lax.scan`.
  
      - `buffers`: If passed, then every leaf of `tree_leaves(buffers(init))` must
          be an array; all such arrays become buffers supporting only `[]` and
          `.at[].set()`. However they will act efficiently, without spurious copies.
          You should avoid performing in-place updates to any quantity that is not a
          buffer.
  
      - `kind`: The type of scan that is lowered to underneath. This may either be
          `"lax"` or `"checkpointed"`.
  
          If `kind` is `"lax"` then the usual `lax.scan` is used.
  
          If `kind` is `"checkpointed"` then the scan uses checkpointing to reduce memory
          usage. It will not be forward-mode autodifferentiable.
  
      - `checkpoints`: Only used if `kind="checkpointed"`. Specifies the number of
          checkpoints to use; if `None` then this is set proportional to `sqrt(length)`.
          Can also be a string `"all"`, representing checkpointing every step.
  
      !!! Danger
  
          Note that `buffers` is subject to the same restrictions as
          `equinox.internal.while_loop`.
  
      Returns:
  
      As `jax.lax.scan`.
      """
  
      init, xs = jtu.tree_map(jnp.asarray, (init, xs))
  
      if kind == "lax":
          return lax.scan(f, init, xs, length)
  
      lengths = {jnp.shape(x)[0] for x in jtu.tree_leaves(xs)}
      if length is not None:
          lengths.add(length)
      if len(lengths) == 1:
          length = lengths.pop()
      else:
          raise ValueError(
              f"Got inconsistent lengths in scan: {lengths}."
          )
      del lengths
  
      if checkpoints == "all":
          checkpoints = length
  
      def cond_fun(val):
          return True
  
      def body_fun(val):
          i, carry, ys = val
          x = jtu.tree_map(lambda z: z[i], xs)
          carry, y = f(carry, x)
          ys = jtu.tree_map(lambda z, zs: zs.at[i].set(z), y, ys)
          return i + 1, carry, ys
  
      x0 = jtu.tree_map(lambda z: z[0], xs)
      _, y0_shape = jax.eval_shape(f, init, x0)
      ys = jtu.tree_map(
          lambda z: jnp.empty((length,) + z.shape, z.dtype), y0_shape
      )
      init_val = (0, init, ys)
  
      if buffers is None:
          _buffers = lambda val: val[2]
      else:
          node_or_nodes = buffers(init)
          is_tuple = True
          is_leaf = lambda node: node is node_or_nodes
  
          def _find(node):
              nonlocal is_tuple
              if is_leaf(node):
                  is_tuple = False
  
          jtu.tree_map(_find, init, is_leaf=is_leaf)
          if is_tuple:
              _buffers = lambda val: tuple(buffers(val[1])) + (val[2],)
          else:
              _buffers = lambda val: (buffers(val[1]), val[2])
  
      _, carry, ys = checkpointed_while_loop(
          cond_fun,
          body_fun,
          init_val,
          buffers=_buffers,
          max_steps=length,
          checkpoints=checkpoints,
      )
      return carry, ys
  
file_path: equinox/equinox/internal/_loop/loop.py
imports: from collections.abc import Callable, Sequence
  from typing import Any, Literal, Optional, TypeVar, Union
  import jax
  import jax.lax
  import jax.numpy
  import jax.tree_util
  from jaxtyping import Array, Bool
  from .bounded import bounded_while_loop
  from .checkpointed import checkpointed_while_loop
  from .common import common_rewrite
label_desc: Use image node_id 2 for calling a global function with example usage: scan(f, init, xs, length) and returns: carry,  ys, lax, True, unknown,  carry,  ys