Repository: Lightning-AI/pytorch-lightning
Issue category: Bug Report

Issue:
OptimizerLRScheduler typing does not fit examples
### Bug description

The return type of `LightningModule.configure_optimizers()` is `OptimizerLRScheduler`, see the [source code](https://github.com/Lightning-AI/pytorch-lightning/blob/e6c26d2d22fc4678b2cf47d57697ebae68b09529/src/lightning/pytorch/core/module.py#L954). However, the examples give return types not fitting this return type. See e.g. the example [here](https://github.com/Lightning-AI/pytorch-lightning/blob/e6c26d2d22fc4678b2cf47d57697ebae68b09529/examples/fabric/build_your_own_trainer/run.py#L42-L49).

Furthermore, the `OptimizerLRScheduler` is only used as a return type, but I don't see where it is actually used, i.e. the other part of the typed interface. A [search for it](https://github.com/search?q=repo%3ALightning-AI%2Fpytorch-lightning%20OptimizerLRScheduler&type=code) does not reveal it.



### What version are you seeing the problem on?

2.3.1

### How to reproduce the bug

Just run mypy on an example, e.g. https://github.com/Lightning-AI/pytorch-lightning/blob/e6c26d2d22fc4678b2cf47d57697ebae68b09529/examples/fabric/build_your_own_trainer/run.py#L42-L49.

### Error messages and logs

Running mypy on this causes a bug:
```
Error: Incompatible return value type (got "tuple[Adam, dict[str, object]]", expected "Optimizer | Sequence[Optimizer] | tuple[Sequence[Optimizer], Sequence[LRScheduler | ReduceLROnPlateau | LRSchedulerConfig]] | OptimizerLRSchedulerConfig | Sequence[OptimizerLRSchedulerConfig] | None")
```


### Environment

<details>
  <summary>Current environment</summary>

* CUDA:
        - GPU:               None
        - available:         False
        - version:           None
* Lightning:
        - lightning-utilities: 0.11.3.post0
        - pytorch-lightning: 2.3.1
        - torch:             2.3.1
        - torchmetrics:      0.8.0
        - torchvision:       0.18.1
* Packages:
        - absl-py:           2.1.0
        - aenum:             3.1.15
        - aiohttp:           3.9.5
        - aiosignal:         1.3.1
        - alabaster:         0.7.16
        - albumentations:    1.3.1
        - antlr4-python3-runtime: 4.9.3
        - arabic-reshaper:   3.0.0
        - asn1crypto:        1.5.1
        - async-timeout:     4.0.3
        - attrs:             23.2.0
        - babel:             2.15.0
        - boto3:             1.34.137
        - botocore:          1.34.137
        - build:             1.2.1
        - certifi:           2024.6.2
        - cffi:              1.16.0
        - charset-normalizer: 3.3.2
        - click:             8.1.7
        - coloredlogs:       15.0.1
        - contourpy:         1.2.1
        - coverage:          5.3.1
        - cryptography:      42.0.8
        - cssselect2:        0.7.0
        - cycler:            0.12.1
        - data-gradients:    0.3.2
        - deprecated:        1.2.14
        - docutils:          0.17.1
        - einops:            0.3.2
        - exceptiongroup:    1.2.1
        - filelock:          3.15.4
        - flatbuffers:       24.3.25
        - fonttools:         4.53.0
        - frozenlist:        1.4.1
        - fsspec:            2024.6.1
        - future:            1.0.0
        - grpcio:            1.64.1
        - html5lib:          1.1
        - huggingface-hub:   0.23.4
        - humanfriendly:     10.0
        - hydra-core:        1.3.2
        - idna:              3.7
        - imagededup:        0.3.1
        - imageio:           2.34.2
        - imagesize:         1.4.1
        - iniconfig:         2.0.0
        - jinja2:            3.1.4
        - jmespath:          1.0.1
        - joblib:            1.4.2
        - json-tricks:       3.16.1
        - jsonschema:        4.22.0
        - jsonschema-specifications: 2023.12.1
        - kiwisolver:        1.4.5
        - lazy-loader:       0.4
        - lightly:           1.5.8
        - lightly-train:     0.1.0
        - lightly-utils:     0.0.2
        - lightning-utilities: 0.11.3.post0
        - lxml:              5.2.2
        - markdown:          3.6
        - markdown-it-py:    3.0.0
        - markupsafe:        2.1.5
        - matplotlib:        3.9.0
        - mdurl:             0.1.2
        - mpmath:            1.3.0
        - multidict:         6.0.5
        - mypy:              1.10.1
        - mypy-extensions:   1.0.0
        - networkx:          3.3
        - numpy:             1.23.0
        - omegaconf:         2.3.0
        - onnx:              1.15.0
        - onnxruntime:       1.15.0
        - onnxsim:           0.4.36
        - opencv-python:     4.10.0.84
        - opencv-python-headless: 4.10.0.84
        - oscrypto:          1.3.0
        - packaging:         24.1
        - pandas:            2.2.2
        - pillow:            10.4.0
        - pip:               24.1.1
        - pip-tools:         7.4.1
        - platformdirs:      4.2.2
        - pluggy:            1.5.0
        - protobuf:          3.20.3
        - psutil:            6.0.0
        - pycparser:         2.22
        - pydantic:          1.10.17
        - pydeprecate:       0.3.2
        - pygments:          2.18.0
        - pyhanko:           0.25.0
        - pyhanko-certvalidator: 0.26.3
        - pyparsing:         3.1.2
        - pypdf:             4.2.0
        - pypng:             0.20220715.0
        - pyproject-hooks:   1.1.0
        - pytest:            8.2.2
        - pytest-mock:       3.14.0
        - python-bidi:       0.4.2
        - python-dateutil:   2.9.0.post0
        - pytorch-lightning: 2.3.1
        - pytz:              2024.1
        - pywavelets:        1.6.0
        - pyyaml:            6.0.1
        - qrcode:            7.4.2
        - qudida:            0.0.4
        - rapidfuzz:         3.9.3
        - referencing:       0.35.1
        - reportlab:         3.6.13
        - requests:          2.32.3
        - rich:              13.7.1
        - rpds-py:           0.18.1
        - ruff:              0.5.0
        - s3transfer:        0.10.2
        - safetensors:       0.4.3
        - scikit-image:      0.24.0
        - scikit-learn:      1.5.0
        - scipy:             1.13.1
        - seaborn:           0.13.2
        - selftrain:         0.1.0
        - setuptools:        70.2.0
        - six:               1.16.0
        - snowballstemmer:   2.2.0
        - sphinx:            4.0.3
        - sphinx-rtd-theme:  1.3.0
        - sphinxcontrib-applehelp: 1.0.8
        - sphinxcontrib-devhelp: 1.0.6
        - sphinxcontrib-htmlhelp: 2.0.5
        - sphinxcontrib-jquery: 4.1
        - sphinxcontrib-jsmath: 1.0.1
        - sphinxcontrib-qthelp: 1.0.7
        - sphinxcontrib-serializinghtml: 1.1.10
        - stringcase:        1.2.0
        - super-gradients:   3.7.1
        - svglib:            1.5.1
        - sympy:             1.12.1
        - tensorboard:       2.17.0
        - tensorboard-data-server: 0.7.2
        - termcolor:         1.1.0
        - threadpoolctl:     3.5.0
        - tifffile:          2024.6.18
        - timm:              1.0.7
        - tinycss2:          1.3.0
        - tomli:             2.0.1
        - torch:             2.3.1
        - torchmetrics:      0.8.0
        - torchvision:       0.18.1
        - tqdm:              4.66.4
        - treelib:           1.6.1
        - typing-extensions: 4.12.2
        - tzdata:            2024.1
        - tzlocal:           5.2
        - uritools:          4.0.3
        - urllib3:           2.2.2
        - webencodings:      0.5.1
        - werkzeug:          3.0.3
        - wheel:             0.43.0
        - wrapt:             1.16.0
        - xhtml2pdf:         0.2.11
        - yarl:              1.9.4
* System:
        - OS:                Darwin
        - architecture:
                - 64bit
                - 
        - processor:         arm
        - python:            3.10.8
        - release:           23.5.0
        - version:           Darwin Kernel Version 23.5.0: Wed May  1 20:12:58 PDT 2024; root:xnu-10063.121.3~5/RELEASE_ARM64_T6000

### More info

_No response_


Hints:
Hey @MalteEbner 
The optimizer should be inside the dict in that example. Thank you for noticing it. Would you like to send a PR with this quick fix? Would be much appreciated 😃 
> Hey @MalteEbner The optimizer should be inside the dict in that example. Thank you for noticing it. Would you like to send a PR with this quick fix? Would be much appreciated 😃

Unfortunately, the problem is not that the example is wrong. This is just a symptom of an underlying problem. The problem is that the `LightningModule.configure_optimizers() -> OptimizerLRScheduler` and its usage don't fit. 

Toe see why, have a look at where `configure_optimizers()` is used, see the source code [here](https://github.com/Lightning-AI/pytorch-lightning/blob/f91349c961103af48091654775248789b6e03bd1/src/lightning/pytorch/core/optimizer.py#L179-L200): The output of it is stored under the variable name `optim_conf` and then passed to the function
```python
def _configure_optimizers(
    optim_conf: Union[Dict[str, Any], List, Optimizer, Tuple],
)
```

However, these definition type `OptimizerLRScheduler` and the usage type `Union[Dict[str, Any], List, Optimizer, Tuple]  `don't align. To fix this, more is needed

1. Change it to `optim_conf: OptimizerLRScheduler`, so that the usage of `configure_optimizers()` has the same type as its definition.
2. Redefine the `OptimizerLRScheduler` such that it fits the supported types in `_configure_optimizers`. 
I'm not seeing any mypy errors regarding this. Which version did you use and what was the command you ran? The version we test with you can see here: https://github.com/Lightning-AI/pytorch-lightning/blob/master/requirements/typing.txt

We typically bump it together when an new torch version comes out (which is soon again). Maybe your issue will show up, but I'm not seeing it locally. 

Yes sure, the internal `_configure_optimizers` uses a bit more generic typing. Feel free to update it to be more specific 👍 . The return type of `LightningModule.configure_optimizers()` should not be changed, this looks all good to me still. 

Patch:
diff --git a/examples/fabric/build_your_own_trainer/run.py b/examples/fabric/build_your_own_trainer/run.py
index 01044f5d94fa8..c0c2ff28ddc41 100644
--- a/examples/fabric/build_your_own_trainer/run.py
+++ b/examples/fabric/build_your_own_trainer/run.py
@@ -41,7 +41,8 @@ def training_step(self, batch, batch_idx: int):
 
     def configure_optimizers(self):
         optim = torch.optim.Adam(self.parameters(), lr=1e-4)
-        return optim, {
+        return {
+            "optimizer": optim,
             "scheduler": torch.optim.lr_scheduler.ReduceLROnPlateau(optim, mode="max", verbose=True),
             "monitor": "val_accuracy",
             "interval": "epoch",


Test patch: