code: def checkpoint(
      target: Target,
      variables: lift.CollectionFilter = True,
      rngs: lift.PRNGSequenceFilter = True,
      concrete: bool = False,
      prevent_cse: bool = True,
      static_argnums: Union[int, Tuple[int, ...]] = (),
      policy: Optional[Callable[..., bool]] = None,
      methods=None,
  ) -> Target:
      """Lifted version of ``jax.checkpoint``.
  
      Checkpointing is a technique for reducing memory usage by recomputing
      activations during backpropagation. When training large models, it can be
      helpful to checkpoint parts of the model to trade off memory usage for
      additional computation.
  
      Example::
  
        >>> import jax
        >>> import jax.numpy as jnp
        >>> import flax.linen as nn
        ...
        >>> class CheckpointedMLP(nn.Module):
        ...   @nn.compact
        ...   def __call__(self, x):
        ...     CheckpointDense = nn.checkpoint(nn.Dense)
        ...     x = CheckpointDense(128)(x)
        ...     x = nn.relu(x)
        ...     x = CheckpointDense(1)(x)
        ...     return x
        ...
        >>> model = CheckpointedMLP()
        >>> variables = model.init(jax.random.key(0), jnp.ones((1, 16)))
  
      This function is aliased to ``remat`` just like ``jax.remat``.
  
      Args:
        target: a ``Module`` or a function taking a ``Module``
          as its first argument. intermediate computations will be
          re-computed when computing gradients for the target.
        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`` should 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``.
        methods: An optional list of method names that will be lifted, if ``methods``
          is None (default) only the ``__call__`` method will be lifted. If``target``
          is a function, ``methods`` is ignored.
  
      Returns:
        A wrapped version of ``target``. When computing gradients intermediate
        computations will be re-computed on the backward pass.
      """
      # subtract 1 from each static_argnums because 'self' is not passed to the
      # lifted function
      static_argnums = jax.tree_util.tree_map(
          lambda x: x - 1, static_argnums
      )
      return lift_transform(
          lift.checkpoint,
          target,
          variables=variables,
          rngs=rngs,
          concrete=concrete,
          static_argnums=static_argnums,
          prevent_cse=prevent_cse,
          policy=policy,
          methods=methods,
      )
  
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 13 for calling a global function with example usage: checkpoint(target, variables, rngs, concrete, prevent_cse, static_argnums, policy, methods) and returns: lift_transform