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","`","."]
docstr: Set up an optimizer for a model wrapped with XLAFSDP.
  
  This setup method doesn't modify the optimizer or wrap the optimizer. The only thing it currently does is verify
  that the optimizer was created after the model was wrapped with :meth:`setup_module` with a reference to the
  flattened parameters.
docstr_tok: ["Set","up","an","optimizer","for","a","model","wrapped","with","XLAFSDP",".","This","setup","method","does","n't","modify","the","optimizer","or","wrap","the","optimizer",".","The","only","thing","it","currently","does","is","verify","that","the","optimizer","was","created","after","the","model","was","wrapped","with",":","meth",":","`","setup_module","`","with","a","reference","to","the","flattened","parameters","."]
code: def setup_optimizer(self, optimizer: Optimizer) -> Optimizer:
      """Set up an optimizer for a model wrapped with XLAFSDP.
  
      This setup method doesn't modify the optimizer or wrap the optimizer. The only thing it currently does is verify
      that the optimizer was created after the model was wrapped with :meth:`setup_module` with a reference to the
      flattened parameters.
  
      """
      if any(
          getattr(p, "_is_sharded", False)
          for group in optimizer.param_groups
          for p in group["params"]
      ):
          return optimizer
      raise ValueError(
          "The optimizer does not seem to reference any XLAFSDP parameters. HINT: Make sure to create the optimizer"
          " after setting up the model."
      )
  
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 20 for calling the XLAFSDPStrategy obj's underlying member method code with example usage: obj.setup_optimizer(optimizer) and returns: optimizer