class_docstr: Strategy for training multiple TPU devices using the :func:`torch_xla.distributed.xla_multiprocessing.spawn`
  method.
class_docstr_tok: ["Strategy","for","training","multiple","TPU","devices","using","the",":","func",":","`","torch_xla.distributed.xla_multiprocessing.spawn","`","method","."]
code: def setup(self, trainer: "pl.Trainer") -> None:
      assert self.accelerator
      self.accelerator.setup(trainer)
  
      if self.debug:
          os.environ["PT_XLA_DEBUG"] = "1"
  
      assert self.lightning_module
      shared_params = find_shared_parameters(self.lightning_module)
      self.model_to_device()
  
      set_shared_parameters(self.lightning_module, shared_params)
      self.setup_precision_plugin()
  
      if self._sync_module_states:
          if _XLA_GREATER_EQUAL_2_1:
              from torch_xla.core.xla_model import (
                  broadcast_master_param,
              )
          else:
              from torch_xla.experimental.pjrt import (
                  broadcast_master_param,
              )
  
          broadcast_master_param(self.model)
  
      if trainer.state.fn == TrainerFn.FITTING:
          self.setup_optimizers(trainer)
          _optimizers_to_device(self.optimizers, self.root_device)
  
imports: import io
  import os
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
  import torch
  from torch import Tensor
  from torch.nn import Module
  from typing_extensions import override
  import lightning.pytorch
  from lightning.fabric.accelerators.xla import _XLA_AVAILABLE, _XLA_GREATER_EQUAL_2_1
  from lightning.fabric.plugins import XLACheckpointIO
  from lightning.fabric.plugins.environments import XLAEnvironment
  from lightning.fabric.strategies import _StrategyRegistry
  from lightning.fabric.utilities.optimizer import _optimizers_to_device
  from lightning.fabric.utilities.types import _PATH, ReduceOp
  from lightning.pytorch.plugins import XLAPrecision
  from lightning.pytorch.plugins.io.wrapper import _WrappingCheckpointIO
  from lightning.pytorch.strategies.ddp import DDPStrategy
  from lightning.pytorch.strategies.launchers.xla import _XLALauncher
  from lightning.pytorch.strategies.strategy import TBroadcast
  from lightning.pytorch.trainer.states import TrainerFn
  from lightning.pytorch.utilities import find_shared_parameters, set_shared_parameters
  from lightning.pytorch.utilities.rank_zero import rank_zero_only
label_desc: Use image node_id 12 for calling the XLAStrategy obj's underlying member method code with example usage: obj.setup(trainer) without return types