code: def remat_scan(
      target: Target,
      lengths: Optional[Sequence[int]] = (),
      policy: Optional[Callable[..., bool]] = None,
      variable_broadcast: lift.CollectionFilter = False,
      variable_carry: lift.CollectionFilter = False,
      variable_axes: Mapping[
          lift.CollectionFilter, lift.InOutScanAxis
      ] = FrozenDict({True: 0}),
      split_rngs: Mapping[lift.PRNGSequenceFilter, bool] = FrozenDict(
          {True: True}
      ),
  ) -> Target:
      """Combines remat and scan for memory efficiency and constant time compilation.
  
      ``remat_scan`` allows for constant compile times and sublinear
      memory usage with respect to model depth. At a small constant
      penalty. This is typically beneficial for very deep models.
  
      Example::
  
        >>> import flax.linen as nn
  
        >>> class BigModel(nn.Module):
        ...   @nn.compact
        ...   def __call__(self, x):
        ...     DenseStack = nn.remat_scan(nn.Dense, lengths=(10, 10))
        ...     # 100x dense with O(sqrt(N)) memory for gradient computation
        ...     return DenseStack(8, name="dense_stack")(x)
  
      Args:
        target: a ``Module`` or a function taking a ``Module`` as its first
          argument.
        lengths: number of loop iterations at the given level. The total number of
          iterations ``n = prod(lengths)``. each loop is rematerialized. This way the
          memory consumption is proportional to ``n^(1 / d)`` where ``d =
          len(lengths)``. Minimal memory consumptions requires tuning the lengths
          such that the same amount of memory is consumed at each level of the
          nested loop.
        policy: Experimental checkpoint policy, see ``jax.checkpoint``.
        variable_broadcast: Specifies the broadcasted variable collections. A
          broadcasted variable should not depend on any computation that cannot be
          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.
        variable_axes: the variable collections that are scanned over. Defaults to
          ``{True: 0}``.
        split_rngs: Split PRNG sequences will be different for each loop iterations.
          If split is False the PRNGs will be the same across iterations. Defaults
          to ``{True: True}``.
  
      Returns:
        A wrapped version of ``target`` that repeats itself prod(lengths) times.
      """
      return lift_transform(
          lift.remat_scan,
          target,
          lengths=lengths,
          variable_broadcast=variable_broadcast,
          variable_carry=variable_carry,
          variable_axes=variable_axes,
          split_rngs=split_rngs,
          policy=policy,
      )
  
imports: import dataclasses
  import functools
  import inspect
  from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union
  from flax import errors, struct, traceback_util
  from flax.core import Scope, lift, meta
  from flax.core.frozen_dict import FrozenDict
  from flax.linen import module
  from flax.linen.module import Module, Variable, _derive_profiling_name, _get_unbound_fn, wrap_method_once
  import jax
label_desc: Use image node_id 14 for calling a global function with example usage: remat_scan(target, lengths, policy, variable_broadcast, variable_carry, variable_axes, split_rngs) and returns: lift_transform