code: def selective_checkpoint_wrapper(
      module: torch.nn.Module,
      memory_budget: Optional[float] = None,
      policy_fn: Optional[Callable] = None,
  ):
      """
      Wrap a module with selective activation checkpointing.
  
      It behaves similarly to PyTorch's checkpoint_wrapper, but gives the possibility
      to the user to either specify a handcrafted policy_fn, or to let an optimization
      algorithm to select the policy given a user-specified memory_budget.
  
      The user should either specify the memory_budget argument or the policy_fn.
  
      The memory_budget is a float value between 0 (recompute everything in the backward) or 1
      (store everything for backward). Using a value of 0 should be similar to PyTorch's
      activation checkpoint, while 1 should be similar to the behavior of not using any
      activation checkpointing.
      """
      return SelectiveCheckpointWrapper(
          module, memory_budget, policy_fn
      )
  
imports: import time
  import weakref
  from collections import defaultdict
  from contextlib import nullcontext
  from copy import deepcopy
  from dataclasses import astuple, dataclass
  from typing import Any, Callable, ContextManager, Dict, List, Optional, Tuple
  import torch
  from torch.testing._internal.composite_compliance import is_inplace, is_inplace_view_fn, is_view_fn
  from torch.utils._python_dispatch import TorchDispatchMode
  from torch.utils._pytree import tree_map
  from torch.utils.checkpoint import get_device_states, set_device_states
label_desc: Use image node_id 8 for calling a global function with example usage: selective_checkpoint_wrapper(module, memory_budget, policy_fn) and returns: SelectiveCheckpointWrapper