Incompatible typing according to mypy

[DanielPerezJensen]

Hi there, since the introduction of Typing in Python I have tried to incorporate this into my workflow. I wanted to type one of my PyTorch Lightning models by typing the various forward and step functions of the model. I believe I have typed it correctly, but running mypy (v. 0.950) gives the following errors and warnings:

models.py:33: error: Signature of "forward" incompatible with supertype "LightningModule"
models.py:33: note:      Superclass:
models.py:33: note:          def forward(self, *args: Any, **kwargs: Any) -> Any
models.py:33: note:      Subclass:
models.py:33: note:          def forward(self, x_dict: Dict[str, Tensor], edge_index_dict: Dict[str, Tensor], h_dict: Optional[Dict[str, Tensor]] = ..., c_dict: Optional[Dict[str, Tensor]] = ...) -> Dict[str, Tensor]
models.py:52: error: Signature of "training_step" incompatible with supertype "LightningModule"
models.py:52: note:      Superclass:
models.py:52: note:          def training_step(self, *args: Any, **kwargs: Any) -> Union[Tensor, Dict[str, Any]]
models.py:52: note:      Subclass:
models.py:52: note:          def training_step(self, batch: Any, batch_idx: int) -> Tensor
models.py:63: error: Signature of "validation_step" incompatible with supertype "LightningModule"
models.py:63: note:      Superclass:
models.py:63: note:          def validation_step(self, *args: Any, **kwargs: Any) -> Union[Tensor, Dict[str, Any], None]
models.py:63: note:      Subclass:
models.py:63: note:          def validation_step(self, batch: Any, batch_idx: int) -> Tensor
models.py:71: error: Signature of "test_step" incompatible with supertype "LightningModule"
models.py:71: note:      Superclass:
models.py:71: note:          def test_step(self, *args: Any, **kwargs: Any) -> Union[Tensor, Dict[str, Any], None]
models.py:71: note:      Subclass:
models.py:71: note:          def test_step(self, batch: Any, batch_idx: int) -> Tensor
Found 4 errors in 1 file (checked 1 source file)

I was just wondering if this is due to me mistyping, which I am pretty sure I haven't, or that it might be some issue in MyPy or PyTorch Lightning. Any guidance is appreciated. The script I used is shown below, not sure if you can import it as it is dependant on some other files I have but if you need those as well I can add it down the line.

import torch
import pytorch_lightning as pl
from torch import nn
from torch_geometric.data.batch import Batch
from base_models import HeteroGLSTM
from omegaconf import DictConfig, OmegaConf

from typing import Any, Optional, Tuple, List, Dict, Callable

# Typing options
TensorDict = Dict[str, torch.Tensor]


class HeteroGLSTM_pl(pl.LightningModule):
    def __init__(
        self,
        cfg: DictConfig,
        metadata: Tuple[List[str], List[Tuple[str, str, str]]],
    ) -> None:
        super().__init__()

        self.model = HeteroGLSTM(cfg.model.num_layers,
                                 cfg.model.hidden_channels,
                                 cfg.model.n_outputs, metadata)
        self.linear = nn.Linear(cfg.model.hidden_channels, cfg.model.n_outputs)

        self.activation = nn.ReLU()
        self.loss_fn = nn.MSELoss()

        # Store some hyperparameters
        self.cfg = cfg

    def forward(
        self, x_dict: TensorDict, edge_index_dict: TensorDict,
        h_dict: Optional[TensorDict] = None,
        c_dict: Optional[TensorDict] = None
    ) -> TensorDict:

        h, c = self.model(x_dict, edge_index_dict,
                          h_dict=h_dict, c_dict=c_dict)
        out = self.activation(h["measurement"])

        out = self.linear(out)

        return out

    def configure_optimizers(self) -> torch.optim.Optimizer:
        # Extract optimizer
        self.optimizer = get_optimizer(self.cfg.optimizer.name)
        return self.optimizer(self.parameters(), self.cfg.optimizer.lr)

    def training_step(self, batch: Batch, batch_idx: int) -> torch.Tensor:
        out = self(batch.x_dict, batch.edge_index_dict)
        out = out.squeeze()

        loss = self.loss_fn(batch.y_dict["measurement"], out)

        self.log("train_loss", loss, on_epoch=True, on_step=True,
                 batch_size=self.cfg.training.batch_size)

        return loss

    def validation_step(self, batch: Batch, batch_idx: int) -> torch.Tensor:
        loss = self._shared_eval_step(batch, batch_idx)

        self.log("val_loss", loss, on_epoch=True, on_step=True,
                 batch_size=self.cfg.training.batch_size)

        return loss

    def test_step(self, batch: Batch, batch_idx: int) -> torch.Tensor:
        loss = self._shared_eval_step(batch, batch_idx)

        self.log("test_loss", loss,
                 batch_size=self.cfg.training.batch_size)
        return loss

    def _shared_eval_step(self, batch: Batch, batch_idx: int) -> torch.Tensor:
        out = self(batch.x_dict, batch.edge_index_dict)
        out = out.squeeze()

        loss = self.loss_fn(batch.y_dict["measurement"], out)

        return loss


def get_optimizer(optimizer_name: str) -> Callable[..., torch.optim.Optimizer]:

    optimizer_dict = {
        "adam": torch.optim.Adam,
        "adamw": torch.optim.AdamW
    }

    return optimizer_dict[optimizer_name]

[DanielPerezJensen]

Commenting this to bump it :)

[akihironitta]

@DanielPerezJensen I think the cause is that you're overriding their typing as the error says. For example, in your script, training_step has torch.Tensor as its return type, but that of the super class (LightningModule) has Union[Tensor, Dict[str, Any]].

If you're intentionally overriding the type, you can just mark it as type: ignore[override].

    def training_step(self, batch: Batch, batch_idx: int) -> torch.Tensor:  # type: ignore[override]

[DanielPerezJensen]

I can't really see how else I am supposed to type the function. In the error you show above, it says that the args can be Any, and that the output is Union[Tensor, Dict[str, Any]]. Which I seem to have done right? There is no wrong type here.

[DanielPerezJensen]

When I for example change it to this:

    def training_step(
        self, batch: Batch, batch_idx: int
    ) -> Union[Tensor, Dict[str, Tensor]]:

It still throws an error. Not really sure what is going on ;p

[DanielPerezJensen]

How do I go about solving this then?

[akihironitta]

@DanielPerezJensen You need to update the args/kwargs to make it compatible with the super class. Let me remind you again that you can ignore errors if you need. I'm sorry but I don't think there's anything I can help with it further here.

[DanielPerezJensen]

Thanks for the help thus far, but I just feel like my typing is correct. It says *args is Any, which should include all typing right? And the output should be Union[Tensor, Dict[str, Any]]. Which again I am fulfilling with Tensor. So I don't understand what exactly it means by Liskov substitution principle.