code: def _auto_wrap_policy_kwargs(
      policy: Optional["_POLICY"], kwargs: Dict
  ) -> Dict:
      if policy is None:
          return kwargs
      if isinstance(policy, set):
          from torch_xla.distributed.fsdp.wrap import (
              transformer_auto_wrap_policy,
          )
  
          # this is not transformer specific despite the name
          policy = partial(
              transformer_auto_wrap_policy, transformer_layer_cls=policy
          )
      kwargs["auto_wrap_policy"] = policy
      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 1 for calling a global function with example usage: _auto_wrap_policy_kwargs(policy, kwargs) and returns: kwargs, kwargs