Update mypy in CI (#19449)

Author: awaelchli

requirements/data/data.txt

@@ -3,6 +3,6 @@
 
 lightning-utilities >=0.8.0, <0.10.0
 # to be able to include also PL 2.0 and preserve `>` needed for CI min version bypass
-torch >0.14.0, <2.2.0
+torch >0.14.0, <=2.2.0
 lightning-cloud
 filelock

requirements/data/examples.txt

@@ -1,3 +1,3 @@
 Pillow >= 9.5.0
 # min version to match torch >= 2.0.1
-torchvision >=0.15.2, <0.17.0
+torchvision >=0.15.2, <0.18.0

requirements/fabric/base.txt

@@ -2,7 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 numpy >=1.17.2, <1.27.0
-torch >=1.13.0, <2.2.0
+torch >=1.13.0, <=2.2.0
 fsspec[http] >=2022.5.0, <2023.11.0
 packaging >=20.0, <=23.1
 typing-extensions >=4.4.0, <4.10.0

requirements/fabric/examples.txt

@@ -1,6 +1,6 @@
 # NOTE: the upper bound for the package version is only set for CI stability, and it is dropped while installing this package
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
-torchvision >=0.14.0, <0.17.0
+torchvision >=0.14.0, <0.18.0
 torchmetrics >=0.10.0, <1.3.0
 lightning-utilities >=0.8.0, <0.10.0

requirements/pytorch/base.txt

@@ -2,7 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 numpy >=1.17.2, <1.27.0
-torch >=1.13.0, <2.2.0
+torch >=1.13.0, <=2.2.0
 tqdm >=4.57.0, <4.67.0
 PyYAML >=5.4, <6.1.0
 fsspec[http] >=2022.5.0, <2023.11.0

requirements/pytorch/examples.txt

@@ -2,7 +2,7 @@
 #  in case you want to preserve/enforce restrictions on the latest compatible version, add "strict" as an in-line comment
 
 requests <2.32.0
-torchvision >=0.14.0, <0.17.0
+torchvision >=0.14.0, <0.18.0
 gym[classic_control] >=0.17.0, <0.27.0
 ipython[all] <8.15.0
 torchmetrics >=0.10.0, <1.3.0

requirements/typing.txt

@@ -1,5 +1,5 @@
-mypy==1.5.1
-torch==2.1.0
+mypy==1.8.0
+torch==2.2.0
 
 types-Markdown
 types-PyYAML

src/lightning/app/cli/lightning_cli.py

@@ -363,7 +363,7 @@ def init_pl_app(source: Union[Tuple[str], Tuple[str, str]], name: str, overwrite
         source_dir = str(Path(script_path).resolve().parent)
     elif len(source) == 2:
         # enable type checking once https://github.com/python/mypy/issues/1178 is available
-        source_dir, script_path = source  # type: ignore
+        source_dir, script_path = source
     else:
         click.echo(
             f"Incorrect number of arguments. You passed ({', '.join(source)}) but only either one argument"

src/lightning/data/streaming/compression.py

@@ -11,7 +11,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from abc import ABC, abstractclassmethod, abstractmethod
+from abc import ABC, abstractmethod
 from typing import Dict, TypeVar
 
 from lightning_utilities.core.imports import RequirementCache, requires
@@ -35,7 +35,8 @@ def compress(self, data: bytes) -> bytes:
     def decompress(self, data: bytes) -> bytes:
         pass
 
-    @abstractclassmethod
+    @classmethod
+    @abstractmethod
     def register(cls, compressors: Dict[str, "Compressor"]) -> None:
         pass
 
@@ -60,7 +61,7 @@ def decompress(self, data: bytes) -> bytes:
         return zstd.decompress(data)
 
     @classmethod
-    def register(cls, compressors: Dict[str, "Compressor"]) -> None:  # type: ignore
+    def register(cls, compressors: Dict[str, "Compressor"]) -> None:
         if not _ZSTD_AVAILABLE:
             return
 

src/lightning/fabric/loggers/csv_logs.py

@@ -135,7 +135,7 @@ def experiment(self) -> "_ExperimentWriter":
 
     @override
     @rank_zero_only
-    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:  # type: ignore[override]
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         raise NotImplementedError("The `CSVLogger` does not yet support logging hyperparameters.")
 
     @override

src/lightning/fabric/loggers/tensorboard.py

@@ -220,7 +220,7 @@ def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None)
 
     @override
     @rank_zero_only
-    def log_hyperparams(  # type: ignore[override]
+    def log_hyperparams(
         self, params: Union[Dict[str, Any], Namespace], metrics: Optional[Dict[str, Any]] = None
     ) -> None:
         """Record hyperparameters. TensorBoard logs with and without saved hyperparameters are incompatible, the

src/lightning/fabric/plugins/precision/amp.py

@@ -90,7 +90,7 @@ def optimizer_step(
         if isinstance(optimizer, LBFGS):
             raise TypeError("AMP and the LBFGS optimizer are not compatible.")
         # note: the scaler will skip the `optimizer.step` if nonfinite gradients are found
-        step_output = self.scaler.step(optimizer, **kwargs)
+        step_output = self.scaler.step(optimizer, **kwargs)  # type: ignore[arg-type]
         self.scaler.update()
         return step_output
 

src/lightning/fabric/plugins/precision/bitsandbytes.py

@@ -344,9 +344,9 @@ def quantize(
         def to_empty(self, *, device: _DEVICE, recurse: bool = True) -> Self:
             if self.weight.dtype == torch.uint8:  # was quantized
                 # cannot init the quantized params directly
-                weight = torch.empty(self.weight.quant_state[1], device=device, dtype=torch.half)  # type: ignore[arg-type]
+                weight = torch.empty(self.weight.quant_state[1], device=device, dtype=torch.half)
             else:
-                weight = torch.empty_like(self.weight.data, device=device)  # type: ignore[arg-type]
+                weight = torch.empty_like(self.weight.data, device=device)
             device = torch.device(device)
             if device.type == "cuda":  # re-quantize
                 self.quantize_(weight, device)

src/lightning/fabric/plugins/precision/fsdp.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import TYPE_CHECKING, Any, ContextManager, Dict, Literal, Optional, cast
+from typing import TYPE_CHECKING, Any, ContextManager, Dict, Literal, Optional
 
 import torch
 from lightning_utilities import apply_to_collection
@@ -128,7 +128,7 @@ def convert_output(self, data: Any) -> Any:
     @override
     def backward(self, tensor: Tensor, model: Optional[Module], *args: Any, **kwargs: Any) -> None:
         if self.scaler is not None:
-            tensor = cast(Tensor, self.scaler.scale(tensor))
+            tensor = self.scaler.scale(tensor)
         super().backward(tensor, model, *args, **kwargs)
 
     @override

src/lightning/fabric/strategies/fsdp.py

@@ -225,7 +225,7 @@ def mixed_precision_config(self) -> Optional["MixedPrecision"]:
             return plugin.mixed_precision_config
         return None
 
-    @property  # type: ignore[override]
+    @property
     @override
     def precision(self) -> FSDPPrecision:
         plugin = self._precision

src/lightning/fabric/strategies/single_xla.py

@@ -50,7 +50,7 @@ def __init__(
             precision=precision,
         )
 
-    @property  # type: ignore[override]
+    @property
     @override
     def checkpoint_io(self) -> XLACheckpointIO:
         plugin = self._checkpoint_io
@@ -66,7 +66,7 @@ def checkpoint_io(self, io: Optional[XLACheckpointIO]) -> None:
             raise TypeError(f"The XLA strategy can only work with the `XLACheckpointIO` plugin, found {io}")
         self._checkpoint_io = io
 
-    @property  # type: ignore[override]
+    @property
     @override
     def precision(self) -> XLAPrecision:
         plugin = self._precision

src/lightning/fabric/strategies/xla.py

@@ -72,7 +72,7 @@ def root_device(self) -> torch.device:
     def num_processes(self) -> int:
         return len(self.parallel_devices) if self.parallel_devices is not None else 0
 
-    @property  # type: ignore[override]
+    @property
     @override
     def checkpoint_io(self) -> XLACheckpointIO:
         plugin = self._checkpoint_io
@@ -88,7 +88,7 @@ def checkpoint_io(self, io: Optional[XLACheckpointIO]) -> None:
             raise TypeError(f"The XLA strategy can only work with the `XLACheckpointIO` plugin, found {io}")
         self._checkpoint_io = io
 
-    @property  # type: ignore[override]
+    @property
     @override
     def precision(self) -> XLAPrecision:
         plugin = self._precision

src/lightning/fabric/strategies/xla_fsdp.py

@@ -124,7 +124,7 @@ def root_device(self) -> torch.device:
     def num_processes(self) -> int:
         return len(self.parallel_devices) if self.parallel_devices is not None else 0
 
-    @property  # type: ignore[override]
+    @property
     @override
     def checkpoint_io(self) -> XLACheckpointIO:
         plugin = self._checkpoint_io
@@ -140,7 +140,7 @@ def checkpoint_io(self, io: Optional[XLACheckpointIO]) -> None:
             raise TypeError(f"The XLA strategy can only work with the `XLACheckpointIO` plugin, found {io}")
         self._checkpoint_io = io
 
-    @property  # type: ignore[override]
+    @property
     @override
     def precision(self) -> XLAPrecision:
         plugin = self._precision

src/lightning/fabric/utilities/apply_func.py

@@ -27,7 +27,7 @@
 
 
 def _from_numpy(value: np.ndarray, device: _DEVICE) -> Tensor:
-    return torch.from_numpy(value).to(device)  # type: ignore[arg-type]
+    return torch.from_numpy(value).to(device)
 
 
 CONVERSION_DTYPES: List[Tuple[Any, Callable[[Any, Any], Tensor]]] = [

src/lightning/fabric/utilities/init.py

@@ -60,7 +60,7 @@ def _materialize(module: torch.nn.Module, device: _DEVICE) -> None:
     """Materialize a module."""
     if not _TORCH_GREATER_EQUAL_2_1:
         raise RuntimeError("recurse=False requires torch 2.1")
-    module.to_empty(device=device, recurse=False)  # type: ignore[arg-type]
+    module.to_empty(device=device, recurse=False)
     if not hasattr(module, "reset_parameters"):
         raise TypeError(
             f"Materialization requires that the `{type(module).__name__}.reset_parameters` method is implemented."

src/lightning/pytorch/cli.py

@@ -360,7 +360,7 @@ def __init__(
         self.trainer_class = trainer_class
         self.trainer_defaults = trainer_defaults or {}
         self.seed_everything_default = seed_everything_default
-        self.parser_kwargs = parser_kwargs or {}  # type: ignore[var-annotated]  # github.com/python/mypy/issues/6463
+        self.parser_kwargs = parser_kwargs or {}
         self.auto_configure_optimizers = auto_configure_optimizers
 
         self.model_class = model_class

src/lightning/pytorch/core/module.py

@@ -320,7 +320,7 @@ def loggers(self) -> Union[List[Logger], List[FabricLogger]]:
             return self._fabric.loggers
         if self._trainer is not None:
             return self._trainer.loggers
-        return []  # type: ignore[return-value]
+        return []
 
     def _call_batch_hook(self, hook_name: str, *args: Any) -> Any:
         trainer = self._trainer

src/lightning/pytorch/loggers/comet.py

@@ -305,7 +305,7 @@ def experiment(self) -> Union["Experiment", "ExistingExperiment", "OfflineExperi
 
     @override
     @rank_zero_only
-    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:  # type: ignore[override]
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         params = _convert_params(params)
         params = _flatten_dict(params)
         self.experiment.log_parameters(params)

src/lightning/pytorch/loggers/csv_logs.py

@@ -145,7 +145,7 @@ def save_dir(self) -> str:
 
     @override
     @rank_zero_only
-    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:  # type: ignore[override]
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         params = _convert_params(params)
         self.experiment.log_hparams(params)
 

src/lightning/pytorch/loggers/mlflow.py

@@ -220,7 +220,7 @@ def experiment_id(self) -> Optional[str]:
 
     @override
     @rank_zero_only
-    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:  # type: ignore[override]
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         params = _convert_params(params)
         params = _flatten_dict(params)
 

src/lightning/pytorch/loggers/neptune.py

@@ -390,7 +390,7 @@ def run(self) -> "Run":
 
     @override
     @rank_zero_only
-    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:  # type: ignore[override]
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         r"""Log hyperparameters to the run.
 
         Hyperparameters will be logged under the "<prefix>/hyperparams" namespace.

src/lightning/pytorch/loggers/tensorboard.py

@@ -154,7 +154,7 @@ def save_dir(self) -> str:
 
     @override
     @rank_zero_only
-    def log_hyperparams(  # type: ignore[override]
+    def log_hyperparams(
         self, params: Union[Dict[str, Any], Namespace], metrics: Optional[Dict[str, Any]] = None
     ) -> None:
         """Record hyperparameters. TensorBoard logs with and without saved hyperparameters are incompatible, the

src/lightning/pytorch/loggers/wandb.py

@@ -415,7 +415,7 @@ def watch(
 
     @override
     @rank_zero_only
-    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:  # type: ignore[override]
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
         params = _convert_params(params)
         params = _sanitize_callable_params(params)
         self.experiment.config.update(params, allow_val_change=True)

src/lightning/pytorch/plugins/precision/amp.py

@@ -86,14 +86,14 @@ def optimizer_step(  # type: ignore[override]
             # Unscaling needs to be performed here in case we are going to apply gradient clipping.
             # Optimizers that perform unscaling in their `.step()` method are not supported (e.g., fused Adam).
             # Note: `unscale` happens after the closure is executed, but before the `on_before_optimizer_step` hook.
-            self.scaler.unscale_(optimizer)
+            self.scaler.unscale_(optimizer)  # type: ignore[arg-type]
 
         self._after_closure(model, optimizer)
 
         # in manual optimization, the closure does not return a value
         if not skip_unscaling:
             # note: the scaler will skip the `optimizer.step` if nonfinite gradients are found
-            step_output = self.scaler.step(optimizer, **kwargs)
+            step_output = self.scaler.step(optimizer, **kwargs)  # type: ignore[arg-type]
             self.scaler.update()
             return step_output
         return closure_result

src/lightning/pytorch/plugins/precision/fsdp.py

@@ -137,7 +137,7 @@ def convert_output(self, data: Any) -> Any:
     @override
     def pre_backward(self, tensor: Tensor, module: "pl.LightningModule") -> Tensor:  # type: ignore[override]
         if self.scaler is not None:
-            tensor = self.scaler.scale(tensor)  # type: ignore[assignment]
+            tensor = self.scaler.scale(tensor)
         return super().pre_backward(tensor, module)
 
     @override

src/lightning/pytorch/strategies/fsdp.py

@@ -213,7 +213,7 @@ def mixed_precision_config(self) -> Optional["MixedPrecision"]:
             return plugin.mixed_precision_config
         return None
 
-    @property  # type: ignore[override]
+    @property
     @override
     def precision_plugin(self) -> FSDPPrecision:
         plugin = self._precision_plugin

src/lightning/pytorch/strategies/single_xla.py

@@ -56,7 +56,7 @@ def __init__(
         )
         self.debug = debug
 
-    @property  # type: ignore[override]
+    @property
     @override
     def checkpoint_io(self) -> Union[XLACheckpointIO, _WrappingCheckpointIO]:
         plugin = self._checkpoint_io
@@ -72,7 +72,7 @@ def checkpoint_io(self, io: Optional[Union[XLACheckpointIO, _WrappingCheckpointI
             raise TypeError(f"The XLA strategy can only work with the `XLACheckpointIO` plugin, found {io}")
         self._checkpoint_io = io
 
-    @property  # type: ignore[override]
+    @property
     @override
     def precision_plugin(self) -> XLAPrecision:
         plugin = self._precision_plugin

src/lightning/pytorch/strategies/xla.py

@@ -70,7 +70,7 @@ def __init__(
         self._launched = False
         self._sync_module_states = sync_module_states
 
-    @property  # type: ignore[override]
+    @property
     @override
     def checkpoint_io(self) -> Union[XLACheckpointIO, _WrappingCheckpointIO]:
         plugin = self._checkpoint_io
@@ -86,7 +86,7 @@ def checkpoint_io(self, io: Optional[Union[XLACheckpointIO, _WrappingCheckpointI
             raise TypeError(f"The XLA strategy can only work with the `XLACheckpointIO` plugin, found {io}")
         self._checkpoint_io = io
 
-    @property  # type: ignore[override]
+    @property
     @override
     def precision_plugin(self) -> XLAPrecision:
         plugin = self._precision_plugin