code: def _activation_checkpointing_auto_wrapper(
      policy: _POLICY_SET, module: Module, *args: Any, **kwargs: Any
  ) -> Module:
      from torch_xla.distributed.fsdp import (
          XlaFullyShardedDataParallel as XLAFSDP,
      )
      from torch_xla.distributed.fsdp import checkpoint_module
  
      module = (
          checkpoint_module(module)
          if isinstance(module, tuple(policy))
          else module
      )
      return XLAFSDP(module, *args, **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 2 for calling a global function with example usage: _activation_checkpointing_auto_wrapper(policy, module) and returns: XLAFSDP