code: def _activation_checkpointing_kwargs(
      policy: Optional[_POLICY_SET], kwargs: Dict
  ) -> Dict:
      if not policy:
          return kwargs
      if "auto_wrapper_callable" in kwargs:
          raise ValueError(
              "You cannot set both `auto_wrapper_callable` and `activation_checkpointing_policy`. Choose one"
          )
      if not isinstance(policy, set):
          raise TypeError(
              f"`activation_checkpointing_policy` must be a set, found {policy}. You can try defining and"
              " passing `auto_wrapper_callable` instead."
          )
      auto_wrapper_callable = partial(
          _activation_checkpointing_auto_wrapper, policy
      )
      kwargs["auto_wrapper_callable"] = auto_wrapper_callable
      return kwargs
  
imports: import io
  from contextlib import ExitStack, nullcontext
  from functools import partial
  from pathlib import Path
  from typing import TYPE_CHECKING, Any, Callable, ContextManager, Dict, List, Literal, Optional, Set, Tuple, Type, Union
  import torch
  from torch import Tensor
  from torch.nn import Module
  from torch.optim import Optimizer
  from torch.utils.data import DataLoader
  from typing_extensions import override
  from lightning.fabric.accelerators import Accelerator
  from lightning.fabric.accelerators.xla import _XLA_AVAILABLE
  from lightning.fabric.plugins import XLAPrecision
  from lightning.fabric.plugins.environments import XLAEnvironment
  from lightning.fabric.plugins.io.xla import XLACheckpointIO
  from lightning.fabric.strategies import ParallelStrategy, _StrategyRegistry
  from lightning.fabric.strategies.fsdp import _apply_filter
  from lightning.fabric.strategies.launchers.xla import _XLALauncher
  from lightning.fabric.strategies.strategy import TBroadcast, _BackwardSyncControl, _Sharded, _validate_keys_for_strict_loading
  from lightning.fabric.utilities.cloud_io import get_filesystem
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_1_13, _TORCH_GREATER_EQUAL_2_0
  from lightning.fabric.utilities.init import _EmptyInit
  from lightning.fabric.utilities.rank_zero import rank_zero_only, rank_zero_warn
  from lightning.fabric.utilities.types import _PATH, Optimizable, ReduceOp
label_desc: Use image node_id 3 for calling a global function with example usage: _activation_checkpointing_kwargs(policy, kwargs) and returns: kwargs, kwargs