simplify torch.Tensor (#16190)

Author: Borda

tests/tests_fabric/test_parity.py

@@ -25,7 +25,7 @@
 from lightning_utilities.core.apply_func import apply_to_collection
 from tests_fabric.helpers.models import RandomDataset
 from tests_fabric.helpers.runif import RunIf
-from torch import nn
+from torch import nn, Tensor
 from torch.nn.parallel.distributed import DistributedDataParallel
 from torch.utils.data import DataLoader
 from torch.utils.data.distributed import DistributedSampler
@@ -131,7 +131,7 @@ def test_boring_lite_model_single_device(precision, strategy, devices, accelerat
         model.load_state_dict(state_dict)
         pure_state_dict = main(lite.to_device, model, train_dataloader, num_epochs=num_epochs)
 
-    state_dict = apply_to_collection(state_dict, torch.Tensor, lite.to_device)
+    state_dict = apply_to_collection(state_dict, Tensor, lite.to_device)
     for w_pure, w_lite in zip(state_dict.values(), lite_state_dict.values()):
         # TODO: This should be torch.equal, but MPS does not yet support this operation (torch 1.12)
         assert not torch.allclose(w_pure, w_lite)

tests/tests_fabric/utilities/test_apply_func.py

@@ -13,14 +13,15 @@
 # limitations under the License.
 import pytest
 import torch
+from torch import Tensor
 
 from lightning_fabric.utilities.apply_func import move_data_to_device
 
 
 @pytest.mark.parametrize("should_return", [False, True])
 def test_wrongly_implemented_transferable_data_type(should_return):
     class TensorObject:
-        def __init__(self, tensor: torch.Tensor, should_return: bool = True):
+        def __init__(self, tensor: Tensor, should_return: bool = True):
             self.tensor = tensor
             self.should_return = should_return
 

tests/tests_fabric/utilities/test_data.py

@@ -4,6 +4,7 @@
 import pytest
 import torch
 from tests_fabric.helpers.models import RandomDataset, RandomIterableDataset
+from torch import Tensor
 from torch.utils.data import BatchSampler, DataLoader, RandomSampler, SequentialSampler
 
 from lightning_fabric.utilities.data import (
@@ -87,7 +88,7 @@ def __init__(self, attribute2, *args, **kwargs):
 
 
 class MyDataLoader(MyBaseDataLoader):
-    def __init__(self, data: torch.Tensor, *args, **kwargs):
+    def __init__(self, data: Tensor, *args, **kwargs):
         self.data = data
         super().__init__(range(data.size(0)), *args, **kwargs)
 
@@ -209,7 +210,7 @@ def test_replace_dunder_methods_dataloader(cls, args, kwargs, arg_names, dataset
 
     for key, value in checked_values.items():
         dataloader_value = getattr(dataloader, key)
-        if isinstance(dataloader_value, torch.Tensor):
+        if isinstance(dataloader_value, Tensor):
             assert dataloader_value is value
         else:
             assert dataloader_value == value
@@ -227,7 +228,7 @@ def test_replace_dunder_methods_dataloader(cls, args, kwargs, arg_names, dataset
 
     for key, value in checked_values.items():
         dataloader_value = getattr(dataloader, key)
-        if isinstance(dataloader_value, torch.Tensor):
+        if isinstance(dataloader_value, Tensor):
             assert dataloader_value is value
         else:
             assert dataloader_value == value

tests/tests_fabric/utilities/test_optimizer.py

@@ -1,6 +1,7 @@
 import collections
 
 import torch
+from torch import Tensor
 
 from lightning_fabric.utilities.optimizer import _optimizer_to_device
 
@@ -22,9 +23,9 @@ def __init__(self, *args, **kwargs):
 def assert_opt_parameters_on_device(opt, device: str):
     for param in opt.state.values():
         # Not sure there are any global tensors in the state dict
-        if isinstance(param, torch.Tensor):
+        if isinstance(param, Tensor):
             assert param.data.device.type == device
         elif isinstance(param, collections.Mapping):
             for subparam in param.values():
-                if isinstance(subparam, torch.Tensor):
+                if isinstance(subparam, Tensor):
                     assert param.data.device.type == device

tests/tests_pytorch/core/test_metric_result_integration.py

@@ -21,6 +21,7 @@
 import torch
 import torchmetrics
 from lightning_utilities.test.warning import no_warning_call
+from torch import Tensor
 from torch.nn import ModuleDict, ModuleList
 from torchmetrics import Metric, MetricCollection
 
@@ -662,7 +663,7 @@ def test_logger_sync_dist(distributed_env, log_val):
     # self.log('bar', 0.5, ..., sync_dist=False)
     meta = _Metadata("foo", "bar")
     meta.sync = _Sync(_should=False)
-    is_tensor = isinstance(log_val, torch.Tensor)
+    is_tensor = isinstance(log_val, Tensor)
 
     if not is_tensor:
         log_val.update(torch.tensor([0, 1]), torch.tensor([0, 0], dtype=torch.long))

tests/tests_pytorch/helpers/datasets.py

@@ -19,6 +19,7 @@
 from typing import Optional, Sequence, Tuple
 
 import torch
+from torch import Tensor
 from torch.utils.data import Dataset
 
 
@@ -69,7 +70,7 @@ def __init__(
         data_file = self.TRAIN_FILE_NAME if self.train else self.TEST_FILE_NAME
         self.data, self.targets = self._try_load(os.path.join(self.cached_folder_path, data_file))
 
-    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
+    def __getitem__(self, idx: int) -> Tuple[Tensor, int]:
         img = self.data[idx].float().unsqueeze(0)
         target = int(self.targets[idx])
 
@@ -125,7 +126,7 @@ def _try_load(path_data, trials: int = 30, delta: float = 1.0):
         return res
 
     @staticmethod
-    def normalize_tensor(tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0) -> torch.Tensor:
+    def normalize_tensor(tensor: Tensor, mean: float = 0.0, std: float = 1.0) -> Tensor:
         mean = torch.as_tensor(mean, dtype=tensor.dtype, device=tensor.device)
         std = torch.as_tensor(std, dtype=tensor.dtype, device=tensor.device)
         return tensor.sub(mean).div(std)
@@ -160,7 +161,7 @@ def __init__(self, root: str, num_samples: int = 100, digits: Optional[Sequence]
         super().__init__(root, normalize=(0.5, 1.0), **kwargs)
 
     @staticmethod
-    def _prepare_subset(full_data: torch.Tensor, full_targets: torch.Tensor, num_samples: int, digits: Sequence):
+    def _prepare_subset(full_data: Tensor, full_targets: Tensor, num_samples: int, digits: Sequence):
         classes = {d: 0 for d in digits}
         indexes = []
         for idx, target in enumerate(full_targets):

tests/tests_pytorch/helpers/deterministic_model.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import torch
-from torch import nn
+from torch import nn, Tensor
 from torch.utils.data import DataLoader, Dataset
 
 from pytorch_lightning.core.module import LightningModule
@@ -56,7 +56,7 @@ def step(self, batch, batch_idx):
 
     def count_num_graphs(self, result, num_graphs=0):
         for k, v in result.items():
-            if isinstance(v, torch.Tensor) and v.grad_fn is not None:
+            if isinstance(v, Tensor) and v.grad_fn is not None:
                 num_graphs += 1
             if isinstance(v, dict):
                 num_graphs += self.count_num_graphs(v)

tests/tests_pytorch/loops/test_evaluation_loop_flow.py

@@ -14,6 +14,7 @@
 """Tests the evaluation loop."""
 
 import torch
+from torch import Tensor
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.core.module import LightningModule
@@ -68,7 +69,7 @@ def backward(self, loss, optimizer, optimizer_idx):
 
     assert len(train_step_out) == 1
     train_step_out = train_step_out[0][0]
-    assert isinstance(train_step_out["loss"], torch.Tensor)
+    assert isinstance(train_step_out["loss"], Tensor)
     assert train_step_out["loss"].item() == 171
 
     # make sure the optimizer closure returns the correct things
@@ -129,7 +130,7 @@ def backward(self, loss, optimizer, optimizer_idx):
 
     assert len(train_step_out) == 1
     train_step_out = train_step_out[0][0]
-    assert isinstance(train_step_out["loss"], torch.Tensor)
+    assert isinstance(train_step_out["loss"], Tensor)
     assert train_step_out["loss"].item() == 171
 
     # make sure the optimizer closure returns the correct things

tests/tests_pytorch/loops/test_training_loop_flow_scalar.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import pytest
-import torch
 from lightning_utilities.test.warning import no_warning_call
+from torch import Tensor
 from torch.utils.data import DataLoader
 from torch.utils.data._utils.collate import default_collate
 
@@ -151,7 +151,7 @@ def backward(self, loss, optimizer, optimizer_idx):
 
     assert len(train_step_out) == 1
     train_step_out = train_step_out[0][0]
-    assert isinstance(train_step_out["loss"], torch.Tensor)
+    assert isinstance(train_step_out["loss"], Tensor)
     assert train_step_out["loss"].item() == 171
 
     # make sure the optimizer closure returns the correct things
@@ -172,7 +172,7 @@ def training_step(self, batch, batch_idx):
             return acc
 
         def training_step_end(self, tr_step_output):
-            assert isinstance(tr_step_output, torch.Tensor)
+            assert isinstance(tr_step_output, Tensor)
             assert self.count_num_graphs({"loss": tr_step_output}) == 1
             self.training_step_end_called = True
             return tr_step_output
@@ -221,7 +221,7 @@ def backward(self, loss, optimizer, optimizer_idx):
 
     assert len(train_step_out) == 1
     train_step_out = train_step_out[0][0]
-    assert isinstance(train_step_out["loss"], torch.Tensor)
+    assert isinstance(train_step_out["loss"], Tensor)
     assert train_step_out["loss"].item() == 171
 
     # make sure the optimizer closure returns the correct things

tests/tests_pytorch/models/test_gpu.py

@@ -183,7 +183,7 @@ def to(self, *args, **kwargs):
 
 @RunIf(min_cuda_gpus=1)
 def test_non_blocking():
-    """Tests that non_blocking=True only gets passed on torch.Tensor.to, but not on other objects."""
+    """Tests that non_blocking=True only gets passed on Tensor.to, but not on other objects."""
     trainer = Trainer()
 
     batch = torch.zeros(2, 3)