code: def scan(
      fn: Callable[..., Any],
      in_axes: Any,
      out_axes: Any,
      length: Optional[int] = None,
      reverse: bool = False,
      unroll: int = 1,
  ):
      """A wrapper around `jax.lax.scan` with in_axes/out_axes api.
  
      Example::
        def body_fn(b, c, x):
          return b + 2, c + 1, 2 * x
  
        loop = scan(body_fn, in_axes=0, out_axes=0)
        broadcast_in = 1
        carry = 2
        xs = jnp.arange(3)
        broadcast_out, carry, ys = loop(broadcast_in, carry, xs)
        print(broadcast_out)  # prints: 3
        print(carry)  # prints: 5
        print(ys)  # prints: [0, 2, 4]
  
  
      Args:
        fn: the body function of the scan loop of the form
          `(broadcast_in, carry, *args) -> (broadcast_out, carry, scan_out)`.
          the broadcast argument allows for loop independent inputs/outputs to
          be computed inside `fn`. `fn` will be called once to compute
          `broadcast_out`. The actual loop will receive `broadcast_out` as the new
          `broadcast_in`. This is useful for initializing values inside the loop.
        in_axes: specifies the axis along which arguments are scanned.
          Use `broadcast` to use the same value across iterations.
        out_axes: specifies the axis along which outputs are concatenated.
          Use `broadcast` if a return value should not be concatenated and
          is independent of the loop body.
        length: number of iterations. Only needs to be specified if there
          is no scan axis from which it can be derived.
        reverse: scan in reverse order from end to start.
        unroll: how many scan iterations to unroll within a single
          iteration of a loop (default: 1).
      Returns:
         the function that performs the scan of the form:
         (broadcast_in, carry_in, *args) -> (broadcast_out, carry_out, scan_out).
      """
  
      def transpose_to_front(ax, xs):
          if ax is broadcast:
              return ()
          if ax == 0:
              return xs
  
          def trans(x):
              perm = tuple(range(x.ndim))
              perm = (ax,) + tuple(np.delete(perm, ax))
              return jnp.transpose(x, perm)
  
          return jax.tree_util.tree_map(trans, xs)
  
      def transpose_from_front(ax, xs):
          if ax is broadcast:
              return ()
          if ax == 0:
              return xs
  
          def trans(x):
              if ax < 0:
                  pax = x.ndim + ax
              else:
                  pax = ax
              assert pax < x.ndim
              perm = (
                  tuple(range(1, pax + 1))
                  + (0,)
                  + tuple(range(pax + 1, x.ndim))
              )
              return jnp.transpose(x, perm)
  
          return jax.tree_util.tree_map(trans, xs)
  
      def scan_fn(broadcast_in, init, *args):
          xs = jax.tree_util.tree_map(transpose_to_front, in_axes, args)
  
          def body_fn(c, xs, init_mode=False):
              # inject constants
              xs = jax.tree_util.tree_map(
                  lambda ax, arg, x: (arg if ax is broadcast else x),
                  in_axes,
                  args,
                  xs,
              )
              broadcast_out, c, ys = fn(broadcast_in, c, *xs)
  
              if init_mode:
                  ys = jax.tree_util.tree_map(
                      lambda ax, y: (y if ax is broadcast else ()),
                      out_axes,
                      ys,
                  )
                  return broadcast_out, ys
              else:
                  ys = jax.tree_util.tree_map(
                      lambda ax, y: (() if ax is broadcast else y),
                      out_axes,
                      ys,
                  )
                  return c, ys
  
          broadcast_body = functools.partial(body_fn, init_mode=True)
  
          carry_avals = jax.tree_util.tree_map(
              lambda x: core.ShapedArray(
                  jnp.shape(x), jnp.result_type(x)
              ),
              init,
          )
          scan_avals = jax.tree_util.tree_map(
              lambda x: core.ShapedArray(
                  jnp.shape(x)[1:], jnp.result_type(x)
              ),
              xs,
          )
          input_avals = (carry_avals, scan_avals)
  
          in_avals, in_tree = jax.tree_util.tree_flatten(input_avals)
          f_flat, out_tree = jax.api_util.flatten_fun_nokwargs(
              lu.wrap_init(broadcast_body), in_tree
          )
          in_pvals = list(map(pe.PartialVal.unknown, in_avals))
          _, out_pvals, _ = pe.trace_to_jaxpr_nounits(f_flat, in_pvals)
  
          out_flat = []
          for pv, const in out_pvals:
              if pv is not None:
                  raise ValueError(
                      "broadcasted variable has a data dependency on the scan body."
                  )
              out_flat.append(const)
          broadcast_in, constants_out = jax.tree_util.tree_unflatten(
              out_tree(), out_flat
          )
  
          c, ys = lax.scan(
              body_fn,
              init,
              xs,
              length=length,
              reverse=reverse,
              unroll=unroll,
          )
          ys = jax.tree_util.tree_map(
              transpose_from_front, out_axes, ys
          )
          ys = jax.tree_util.tree_map(
              lambda ax, const, y: (const if ax is broadcast else y),
              out_axes,
              constants_out,
              ys,
          )
          return broadcast_in, c, ys
  
      return scan_fn
  
file_path: flax/flax/core/axes_scan.py
imports: import functools
  from typing import Any, Callable, Optional
  import jax
  import jax.numpy
  import numpy
  from jax import core, lax
  from jax.extend import linear_util
  from jax.interpreters import partial_eval
label_desc: Use image node_id 1 for calling a global function with example usage: scan(fn, in_axes, out_axes, length, reverse, unroll) and returns: scan_fn, jax, jax, broadcast_in,  c,  ys, , xs, jnp, , xs, jnp, broadcast_out,  ys, c,  ys