code: def checkpoint(
      fn: Callable[..., Any],
      variables: CollectionFilter = True,
      rngs: PRNGSequenceFilter = True,
      concrete: bool = False,
      prevent_cse: bool = True,
      static_argnums: Union[int, Tuple[int, ...]] = (),
      policy: Optional[Callable[..., bool]] = None,
  ) -> Callable[..., Any]:
      """Lifted version of ``jax.checkpoint``.
  
      This function is aliased to ``lift.remat`` just like ``jax.remat``.
  
      Args:
        fn: scope function for which intermediate computations should be
        re-computed when computing gradients.
        variables: The variable collections that are lifted. By default all
          collections are lifted.
        rngs: The PRNG sequences that are lifted. By default all PRNG sequences
          are lifted.
        concrete: Optional, boolean indicating whether ``fun`` may involve
          value-dependent Python control flow (default ``False``). Support for such
          control flow is optional, and disabled by default, because in some
          edge-case compositions with :func:`jax.jit` it can lead to some extra
          computation.
        prevent_cse: Optional, boolean indicating whether to prevent common
          subexpression elimination (CSE) optimizations in the HLO generated from
          differentiation. This CSE prevention has costs because it can foil other
          optimizations, and because it can incur high overheads on some backends,
          especially GPU. The default is True because otherwise, under a ``jit`` or
          ``pmap``, CSE can defeat the purpose of this decorator. But in some
          settings, like when used inside a ``scan``, this CSE prevention mechanism
          is unnecessary, in which case ``prevent_cse`` can be set to False.
        static_argnums: Optional, int or sequence of ints, indicates which argument
          values on which to specialize for tracing and caching purposes. Specifying
          arguments as static can avoid ConcretizationTypeErrors when tracing, but
          at the cost of more retracing overheads.
        policy: Experimental checkpoint policy, see ``jax.checkpoint``.
      Returns:
        A wrapped version of ``fn``. When computing gradients intermediate
        computations will be re-computed when computing gradients.
      """
  
      def inner(
          scope_fn,
          repack_fn,
          variable_groups,
          rng_groups,
          *args,
          **kwargs,
      ):
          # add 2 to each static_argnums because we add two initial arguments to rematted
          static_argnums_ = jax.tree_util.tree_map(
              lambda x: x + 2, static_argnums
          )
  
          @functools.partial(
              jax.remat,
              concrete=concrete,
              static_argnums=static_argnums_,
              prevent_cse=prevent_cse,
              policy=policy,
          )
          @functools.wraps(fn)
          def rematted(variable_groups, rng_groups, *args, **kwargs):
              scope = scope_fn(variable_groups, rng_groups)
              y = fn(scope, *args, **kwargs)
              return y, repack_fn(scope)
  
          return rematted(variable_groups, rng_groups, *args, **kwargs)
  
      return pack(
          inner,
          (variables,),
          (variables,),
          (rngs,),
          name="remat",
          enable_kwargs=True,
      )
  
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 20 for calling a global function with example usage: checkpoint(fn, variables, rngs, concrete, prevent_cse, static_argnums, policy) and returns: pack, rematted, y,  repack_fn