code: def remat_scan(
      body_fn: Callable[..., Any],
      lengths: Sequence[int],
      policy: Optional[Callable[..., bool]] = None,
      variable_broadcast: CollectionFilter = False,
      variable_carry: CollectionFilter = False,
      variable_axes: Mapping[CollectionFilter, InOutScanAxis] = {
          True: 0
      },
      split_rngs: Mapping[PRNGSequenceFilter, bool] = {True: True},
  ) -> Callable[..., Any]:
      """Combines `lift.remat` and `lift.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::
  
        def body_fn(scope, x):
          return nn.dense(scope, x, features=x.shape[-1])
        # 100x dense with O(sqrt(N)) memory for gradient computation
        y = lift.remat_scan(body_fn, lengths=(10, 10))(scope, x)
  
      Args:
        body_fn: Scope function to be repeated using a (nested scan)
        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.
        split_rngs: Split PRNG sequences will be different for each loop iterations.
          If split is False the PRNGs will be the same across iterations.
      Returns:
        A wrapped version of ``body_fn`` that repeats itself prod(lengths) times.
      """
      # TODO(jheek) should remat scan have scan inputs/outputs?
      scan_fn = functools.partial(
          scan,
          variable_broadcast=variable_broadcast,
          variable_carry=variable_carry,
          variable_axes=variable_axes,
          split_rngs=split_rngs,
      )
      if len(lengths) == 1:
  
          def wrapper(scope, carry):
              return body_fn(scope, carry), ()
  
          fn = lambda scope, c: scan_fn(wrapper, length=lengths[0])(
              scope, c
          )[0]
      else:
  
          @functools.partial(remat, policy=policy, prevent_cse=False)
          def inner_loop(scope, carry):
              carry = remat_scan(
                  body_fn,
                  lengths[1:],
                  policy,
                  variable_broadcast,
                  variable_carry,
                  variable_axes,
                  split_rngs,
              )(scope, carry)
              return carry, ()
  
          fn = lambda scope, c: scan_fn(inner_loop, length=lengths[0])(
              scope, c
          )[0]
      return fn
  
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 23 for calling a global function with example usage: remat_scan(body_fn, lengths, policy, variable_broadcast, variable_carry, variable_axes, split_rngs) and returns: fn, body_fn,  , carry,  