code: def scan(
      fn: Callable[..., Any],
      variable_axes: Mapping[CollectionFilter, InOutScanAxis] = {},
      variable_broadcast: CollectionFilter = False,
      variable_carry: CollectionFilter = False,
      split_rngs: Mapping[PRNGSequenceFilter, bool] = {},
      in_axes=0,
      out_axes=0,
      length: Optional[int] = None,
      reverse: bool = False,
      unroll: int = 1,
      data_transform: Optional[Callable[..., Any]] = None,
      metadata_params: Dict[Any, Any] = {},
  ) -> Callable[..., Any]:
      """A lifted version of ``jax.lax.scan``.
  
      See ``jax.lax.scan`` for the unlifted scan in Jax.
  
      To improve consistency with ``vmap``, this version of scan
      uses ``in_axes`` and ``out_axes`` to determine which arguments
      are scanned over and along which axis.
  
      ``scan`` distinguishes between 3 different types of values inside the loop:
  
      1. **scan**: a value that is iterated over in a loop. All scan values must
        have the same size in the axis they are scanned over. Scanned outputs
        will be stacked along the scan axis.
      2. **carry**: A carried value is updated at each loop iteration. It must
        have the same shape and dtype throughout the loop.
      3. **broadcast**: a value that is closed over by the loop. When a variable
        is broadcasted they are typically initialized inside the loop body but
        independent of the loop variables.
  
      The loop body should have the signature
      ``(scope, body, carry, *xs) -> (carry, ys)``, where ``xs`` and ``ys``
      are the scan values that go in and out of the loop.
  
      Example::
  
        scope.variable('counter', 'i', jnp.zeros, ())
        def body_fn(scope, c, x):
          counter = scope.variable('counter', 'i', jnp.zeros, ())
          counter.value += 1
          x = scope.child(nn.dense)(x, 1)
          return c, x
  
        _, ys = lift.scan(
            body_fn,
            variable_carry='counter',
            variable_broadcast='params',
            split_rngs={'params': False})(scope, (), xs)
  
      Args:
        fn: the function to be transformed.
        variable_axes: the variable collections that are scanned over.
        variable_broadcast: Specifies the broadcasted variable collections.
          A broadcasted variable should not depend on any computation that cannot b
          lifted out of the loop. This is typically used to define shared parameters
          inside the fn.
        variable_carry: Specifies the variable collections that are carried through
          the loop. Mutations to these variables are carried to the next iteration
          and will be preserved when the scan finishes.
        split_rngs: Split PRNG sequences will be different for each loop iterations.
          If split is False the PRNGs will be the same across iterations.
        in_axes: Specifies the axis to scan over for the arguments. Should be a
          prefix tree of the arguments. Use `flax.core.broadcast` to feed an entire
          input to each iteration of the scan body.
        out_axes: Specifies the axis to scan over for the return value. Should be a
          prefix tree of the return value.
        length: Specifies the number of loop iterations. This only needs
          to be specified if it cannot be derived from the scan arguments.
        reverse: If true, scan from end to start in reverse order.
        unroll: how many scan iterations to unroll within a single
          iteration of a loop (default: 1).
        data_transform: optional function to transform raw variable and rng groups,
          intended for inline SPMD annotations.
        metadata_params: arguments dict passed to AxisMetadata instances in the
          variable tree.
  
      Returns:
        The scan function with the signature
        ``(scope, carry, *xxs) -> (carry, yys)``, where ``xxs`` and ``yys`` are the
        scan values that go in and out of the loop.
      """
      variable_in_axes, variable_out_axes = _split_in_out_axes(
          variable_axes
      )
      variable_in_groups, variable_in_axes = _unzip2(
          variable_in_axes.items()
      )
      variable_out_groups, variable_out_axes = _unzip2(
          variable_out_axes.items()
      )
      assert all(isinstance(ax, int) for ax in variable_in_axes)
      assert all(isinstance(ax, int) for ax in variable_out_axes)
      rng_groups, rng_splits = _unzip2(split_rngs.items())
      rng_axes = tuple(
          0 if rng_split else axes_scan.broadcast
          for rng_split in rng_splits
      )
  
      def inner(
          scope_fn, repack_fn, variable_groups, rng_groups, init, *args
      ):
          def find_length(axis, x):
              if axis is not axes_scan.broadcast:
                  leaves = jax.tree_util.tree_leaves(x)
                  if leaves:
                      return leaves[0].shape[axis]
              return ()
  
          # split rngs
          lengths = jax.tree_util.tree_map(find_length, in_axes, args)
          lengths = set(jax.tree_util.tree_leaves(lengths))
          if length is None and len(lengths) == 1:
              (d_length,) = lengths
          elif len(lengths) > 1:
              raise ValueError(f"Inconsistent scan lengths: {lengths}")
          elif length is None:
              raise ValueError("length should be specified manually.")
          else:
              d_length = length
          split_fn = lambda rng: random.split(rng, d_length)
  
          rng_groups = tuple(
              tree_map_rngs(split_fn, rng_group) if split else rng_group
              for rng_group, split in zip(rng_groups, rng_splits)
          )
  
          @functools.partial(
              axes_scan.scan,
              in_axes=(variable_in_axes, rng_axes, in_axes),
              out_axes=(out_axes, variable_out_axes),
              length=length,
              reverse=reverse,
              unroll=unroll,
          )
          def scanned(
              broadcast_vars,
              carry,
              scan_variable_groups,
              rng_groups,
              args,
          ):
              carry_vars, c = carry
              variable_groups = (
                  broadcast_vars,
                  carry_vars,
              ) + scan_variable_groups
              if data_transform is not None:
                  variable_groups, rng_groups = data_transform(
                      variable_groups, rng_groups
                  )
              scope = scope_fn(variable_groups, rng_groups)
              c, y = fn(scope, c, *args)
              out_vars = repack_fn(scope)
              broadcast_vars_out = out_vars[0]
              carry_vars = out_vars[1]
              scan_vars = out_vars[2:]
              # add immutable broadcast vars back to broadcast output
              # otherwise they won't be fed to the actual scan body
              for in_group, out_group in zip(
                  broadcast_vars, broadcast_vars_out
              ):
                  for col in in_group:
                      if col not in out_group:
                          out_group[col] = in_group[col]
              return broadcast_vars_out, (carry_vars, c), (y, scan_vars)
  
          broadcast_vars = variable_groups[0]
          carry_vars = variable_groups[1]
          scan_vars = variable_groups[2:]
          new_scan_vars = []
          for scan_group, axis in zip(scan_vars, variable_in_axes):
              new_scan_vars.append(
                  meta.remove_axis(scan_group, axis, metadata_params)
              )
          broadcast_vars, (carry_vars, c), (ys, scan_vars) = scanned(
              broadcast_vars,
              (carry_vars, init),
              tuple(new_scan_vars),
              rng_groups,
              args,
          )
          new_scan_vars = []
          for scan_group, axis in zip(scan_vars, variable_out_axes):
              new_scan_vars.append(
                  meta.add_axis(scan_group, axis, metadata_params)
              )
          scan_vars = tuple(new_scan_vars)
          out_vars = (
              broadcast_vars,
              carry_vars,
          ) + scan_vars
          return (c, ys), out_vars
  
      return pack(
          inner,
          (variable_broadcast, variable_carry) + variable_in_groups,
          (variable_broadcast, variable_carry) + variable_out_groups,
          rng_groups,
          name="scan",
      )
  
imports: import collections
  import dataclasses
  import functools
  from typing import Any, Callable, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Tuple, TypeVar, Union
  import warnings
  from flax import traceback_util
  import jax
  from jax import random
  from .None import axes_scan, meta
  from .frozen_dict import freeze, unfreeze
  from .scope import CollectionFilter, DenyList, Filter, PRNGSequenceFilter, Scope, group_collections, in_filter, intersect_filters, is_filter_empty, subtract_filters, union_filters
label_desc: Use image node_id 15 for calling a global function with example usage: scan(fn, variable_axes, variable_broadcast, variable_carry, split_rngs, in_axes, out_axes, length, reverse, unroll, data_transform, metadata_params) and returns: pack, c,  ys,  out_vars, , broadcast_vars_out,  carry_vars,  c,  y,  scan_vars, leaves