class_docstr: Strategy for training multiple XLA devices using the
  :func:`torch_xla.distributed.xla_fully_sharded_data_parallel.XlaFullyShardedDataParallel` method.
  
  .. warning::  This is an :ref:`experimental <versioning:Experimental API>` feature.
  
  For more information check out https://github.com/pytorch/xla/blob/master/docs/fsdp.md
  
  Args:
      auto_wrap_policy: Same as ``auto_wrap_policy`` parameter in
          :class:`torch_xla.distributed.fsdp.XlaFullyShardedDataParallel`.
          For convenience, this also accepts a set of the layer classes to wrap.
      activation_checkpointing_policy: Used when selecting the modules for
          which you want to enable activation checkpointing. Enabling this can free up a significant amount of memory
          at the cost of speed since activations in these layers need to be recomputed during backpropagation.
          This accepts a set of the layer classes to wrap.
  
      state_dict_type: The format in which the state of the model and optimizers gets saved into the checkpoint.
  
          - ``"full"``: The full weights and optimizer states get assembled on rank 0 and saved to a single file.
          - ``"sharded"``: Each rank saves its shard of weights and optimizer states to a file. The checkpoint is
            a folder with files for each shard in the host. Note that TPU VM multihost does not have a shared
            filesystem.
  
      sequential_save: With this enabled, individual ranks consecutively save their state dictionary shards, reducing
          peak system RAM usage, although it elongates the saving process.
      \**kwargs: See available parameters in :class:`torch_xla.distributed.fsdp.XlaFullyShardedDataParallel`.
class_docstr_tok: ["Strategy","for","training","multiple","XLA","devices","using","the",":","func",":","`","torch_xla.distributed.xla_fully_sharded_data_parallel.XlaFullyShardedDataParallel","`","method",".","..","warning",":",":","This","is","an",":","ref",":","`","experimental","<","versioning",":","Experimental","API",">","`","feature",".","For","more","information","check","out","https",":","\/\/github.com\/pytorch\/xla\/blob\/master\/docs\/fsdp.md","Args",":","auto_wrap_policy",":","Same","as","``","auto_wrap_policy","``","parameter","in",":","class",":","`","torch_xla.distributed.fsdp.XlaFullyShardedDataParallel","`",".","For","convenience",",","this","also","accepts","a","set","of","the","layer","classes","to","wrap",".","activation_checkpointing_policy",":","Used","when","selecting","the","modules","for","which","you","want","to","enable","activation","checkpointing",".","Enabling","this","can","free","up","a","significant","amount","of","memory","at","the","cost","of","speed","since","activations","in","these","layers","need","to","be","recomputed","during","backpropagation",".","This","accepts","a","set","of","the","layer","classes","to","wrap",".","state_dict_type",":","The","format","in","which","the","state","of","the","model","and","optimizers","gets","saved","into","the","checkpoint",".","-","``","``","full","''","``",":","The","full","weights","and","optimizer","states","get","assembled","on","rank","0","and","saved","to","a","single","file",".","-","``","``","sharded","''","``",":","Each","rank","saves","its","shard","of","weights","and","optimizer","states","to","a","file",".","The","checkpoint","is","a","folder","with","files","for","each","shard","in","the","host",".","Note","that","TPU","VM","multihost","does","not","have","a","shared","filesystem",".","sequential_save",":","With","this","enabled",",","individual","ranks","consecutively","save","their","state","dictionary","shards",",","reducing","peak","system","RAM","usage",",","although","it","elongates","the","saving","process",".","\\","*","*","kwargs",":","See","available","parameters","in",":","class",":","`","torch_xla.distributed.fsdp.XlaFullyShardedDataParallel","`","."]
code: def root_device(self) -> torch.device:
      if not self._launched:
          raise RuntimeError(
              "Accessing the XLA device before processes have spawned is not allowed."
          )
      import torch_xla.core.xla_model as xm
  
      return xm.xla_device()
  
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 the XLAFSDPStrategy obj's underlying member method code with example usage: obj.root_device() and returns: xm