code: def scan_trick(
      fn: Callable, intermediates: Sequence[Callable], init: _X
  ) -> _X:
      def body(carry, step):
          out = fn(carry)
          step = nonbatchable(step)
          out = lax.switch(step, intermediates, out)
          return out, None
  
      intermediates = list(intermediates) + [lambda x: x]
      out, _ = lax.scan(body, init, xs=jnp.arange(len(intermediates)))
      return out
  
file_path: equinox/equinox/internal/_misc.py
imports: from collections.abc import Callable, Sequence
  from typing import Any, TypeVar
  import jax.lax
  import jax.numpy
  import jax.tree_util
  from jaxtyping import Array, PyTree, Scalar
  from .._eval_shape import filter_eval_shape
  from .._filters import is_array
  from ._nontraceable import nonbatchable
label_desc: Use image node_id 1 for calling a global function with example usage: scan_trick(fn, intermediates, init) and returns: out, out,  None