Remove Strategy.on_tpu property

CHANGELOG.md

@@ -413,6 +413,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Removed `Strategy.on_gpu` ([#11537](https://github.com/PyTorchLightning/pytorch-lightning/pull/11537))
 
 
+- Removed `Strategy.on_tpu` property ([#11536](https://github.com/PyTorchLightning/pytorch-lightning/pull/11536))
+
+
 ### Fixed
 
 - Fixed security vulnerabilities CVE-2020-1747 and CVE-2020-14343 caused by the `PyYAML` dependency ([#11099](https://github.com/PyTorchLightning/pytorch-lightning/pull/11099))

pytorch_lightning/strategies/parallel.py

@@ -25,7 +25,6 @@
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.strategy import Strategy
-from pytorch_lightning.utilities import _XLA_AVAILABLE
 from pytorch_lightning.utilities.distributed import all_gather_ddp_if_available, ReduceOp
 
 
@@ -49,10 +48,6 @@ def __init__(
     def root_device(self) -> torch.device:
         """Return the root device."""
 
-    @property
-    def on_tpu(self) -> bool:
-        return self.root_device.type == "xla" and _XLA_AVAILABLE
-
     @property
     def lightning_module(self) -> Optional["pl.LightningModule"]:
         return unwrap_lightning_module(self.model) if self.model is not None else None

pytorch_lightning/strategies/single_device.py

@@ -21,7 +21,6 @@
 from pytorch_lightning.plugins.io.checkpoint_plugin import CheckpointIO
 from pytorch_lightning.plugins.precision import PrecisionPlugin
 from pytorch_lightning.strategies.strategy import Strategy
-from pytorch_lightning.utilities import _XLA_AVAILABLE
 from pytorch_lightning.utilities.types import _DEVICE
 
 
@@ -41,10 +40,6 @@ def __init__(
         self.local_rank = 0
         self.world_size = 1
 
-    @property
-    def on_tpu(self) -> bool:
-        return self.root_device.type == "xla" and _XLA_AVAILABLE
-
     def reduce(self, tensor: Any | torch.Tensor, *args: Any, **kwargs: Any) -> Any | torch.Tensor:
         """Reduces a tensor from several distributed processes to one aggregated tensor. As this plugin only
         operates with a single device, the reduction is simply the identity.

pytorch_lightning/strategies/strategy.py

@@ -228,11 +228,6 @@ def batch_to_device(self, batch: Any, device: Optional[torch.device] = None, dat
             return model._apply_batch_transfer_handler(batch, device=device, dataloader_idx=dataloader_idx)
         return move_data_to_device(batch, device)
 
-    @property
-    @abstractmethod
-    def on_tpu(self) -> bool:
-        """Returns whether the current process is done on TPU."""
-
     @property
     @abstractmethod
     def root_device(self) -> torch.device:

tests/strategies/test_ddp_spawn_strategy.py

@@ -57,7 +57,6 @@ def test_ddp_cpu():
     # assert training type plugin attributes for device setting
 
     assert isinstance(trainer.strategy, DDPSpawnStrategy)
-    assert not trainer.strategy.on_tpu
     assert trainer.strategy.root_device == torch.device("cpu")
 
     model = BoringModelDDPCPU()

tests/strategies/test_ddp_strategy.py

@@ -39,7 +39,6 @@ def test_ddp_with_2_gpus():
     trainer = Trainer(gpus=2, strategy="ddp", fast_dev_run=True)
     # assert training type plugin attributes for device setting
     assert isinstance(trainer.strategy, DDPStrategy)
-    assert not trainer.strategy.on_tpu
     local_rank = trainer.strategy.local_rank
     assert trainer.strategy.root_device == torch.device(f"cuda:{local_rank}")
 

tests/strategies/test_single_device_strategy.py

@@ -23,10 +23,9 @@
 
 
 def test_single_cpu():
-    """Tests if on_tpu is set correctly for single CPU strategy."""
+    """Tests if device is set correctly for single CPU strategy."""
     trainer = Trainer()
     assert isinstance(trainer.strategy, SingleDeviceStrategy)
-    assert not trainer.strategy.on_tpu
     assert trainer.strategy.root_device == torch.device("cpu")
 
 
@@ -43,7 +42,6 @@ def test_single_gpu():
     trainer = Trainer(gpus=1, fast_dev_run=True)
     # assert training strategy attributes for device setting
     assert isinstance(trainer.strategy, SingleDeviceStrategy)
-    assert not trainer.strategy.on_tpu
     assert trainer.strategy.root_device == torch.device("cuda:0")
 
     model = BoringModelGPU()

tests/strategies/test_tpu_spawn.py

@@ -96,7 +96,6 @@ def test_model_tpu_one_core():
     trainer = Trainer(tpu_cores=1, fast_dev_run=True, strategy=TPUSpawnStrategy(debug=True))
     # assert training strategy attributes for device setting
     assert isinstance(trainer.strategy, TPUSpawnStrategy)
-    assert trainer.strategy.on_tpu
     assert trainer.strategy.root_device == torch.device("xla", index=1)
     model = BoringModelTPU()
     trainer.fit(model)