Incorrect type hints for LightningModule.save_hyperparameters and hparams

Bug description

Both LightningModule.save_hyperparameters and LightningModule.hparams seem to have type Union[torch._tensor.Tensor, torch.nn.modules.module.Module]. I would expect the former to be Callable and the latter to be an AttributeDict.

How to reproduce the bug

import pytorch_lightning as pl


class LitMNIST(pl.LightningModule):
    def __init__(self, layer_1_dim: int = 128, learning_rate: float = 1e-2):
        super().__init__()
        reveal_type(self.save_hyperparameters)
        self.save_hyperparameters()
        reveal_type(self.hparams)
        self.hparams.layer_1_dim
        self.hparams["learning_rate"]

Error messages and logs

The above code should not have any mypy errors. However, it has several. reveal_type was added to see what type mypy thinks everything should be.

$ mypy --strict test.py
test.py:7: note: Revealed type is "Union[torch._tensor.Tensor, torch.nn.modules.module.Module]"
test.py:8: error: "Tensor" not callable  [operator]
test.py:9: note: Revealed type is "Union[torch._tensor.Tensor, torch.nn.modules.module.Module]"
test.py:10: error: Item "Tensor" of "Union[Tensor, Module]" has no attribute "layer_1_dim"  [union-attr]
test.py:11: error: Value of type "Union[Tensor, Module]" is not indexable  [index]
test.py:11: error: Invalid index type "str" for "Union[Tensor, Module]"; expected type "Union[None, int, slice, Tensor, List[Any], Tuple[Any, ...]]"  [index]
Found 4 errors in 1 file (checked 1 source file)

Environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow): LightningModule
#- PyTorch Lightning Version (e.g., 1.5.0): 1.8.2
#- Lightning App Version (e.g., 0.5.2): N/A
#- PyTorch Version (e.g., 1.10): 1.12.1
#- Python version (e.g., 3.9): 3.10.8
#- OS (e.g., Linux): macOS and Linux
#- CUDA/cuDNN version: N/A
#- GPU models and configuration: N/A
#- How you installed Lightning(`conda`, `pip`, source): Spack
#- Running environment of LightningApp (e.g. local, cloud): local

More info

This issue has been present ever since PyTorch first added public type hints (PyTorch 1.11, March 2022). Since self.hparams is repeatedly accessed throughout every LightningModule, this currently requires type ignores in hundreds of places in TorchGeo.

cc @borda @carmocca

[carmocca]

Here’s the relevant issue upstream: https://github.com/pytorch/pytorch/issues/81462

[adamjstewart]

Thanks, I’ll close this issue and follow that discussion.