feat(fabric): add return-type overloads for Fabric.setup, setup_optimizers and setup_dataloaders

src/lightning/fabric/CHANGELOG.md

@@ -18,7 +18,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
--
+- `Fabric.setup`, `Fabric.setup_optimizers` and `Fabric.setup_dataloaders` now declare return-type overloads, so type checkers infer the concrete wrapped type instead of `Any` ([#22009](https://github.com/Lightning-AI/pytorch-lightning/pull/22009))
 
 ### Removed
 

src/lightning/fabric/fabric.py

@@ -226,6 +226,48 @@ def run(self, *args: Any, **kwargs: Any) -> Any:
 
         """
 
+    @overload
+    def setup(
+        self, module: nn.Module, *, scheduler: None = ..., move_to_device: bool = ..., _reapply_compile: bool = ...
+    ) -> _FabricModule: ...
+
+    @overload
+    def setup(
+        self,
+        module: nn.Module,
+        optimizer: Optimizer,
+        /,
+        *,
+        scheduler: None = ...,
+        move_to_device: bool = ...,
+        _reapply_compile: bool = ...,
+    ) -> tuple[_FabricModule, _FabricOptimizer]: ...
+
+    @overload
+    def setup(
+        self,
+        module: nn.Module,
+        optimizer: Optimizer,
+        /,
+        *,
+        scheduler: "LRScheduler",
+        move_to_device: bool = ...,
+        _reapply_compile: bool = ...,
+    ) -> tuple[_FabricModule, _FabricOptimizer, "LRScheduler"]: ...
+
+    @overload
+    def setup(
+        self,
+        module: nn.Module,
+        optimizer: Optimizer,
+        second_optimizer: Optimizer,
+        /,
+        *optimizers: Optimizer,
+        scheduler: Optional["LRScheduler"] = ...,
+        move_to_device: bool = ...,
+        _reapply_compile: bool = ...,
+    ) -> tuple[Union[_FabricModule, _FabricOptimizer, "LRScheduler"], ...]: ...
+
     def setup(
         self,
         module: nn.Module,
@@ -373,6 +415,14 @@ def setup_module(
         self._models_setup += 1
         return module
 
+    @overload
+    def setup_optimizers(self, optimizer: Optimizer, /) -> _FabricOptimizer: ...
+
+    @overload
+    def setup_optimizers(
+        self, first_optimizer: Optimizer, second_optimizer: Optimizer, /, *optimizers: Optimizer
+    ) -> tuple[_FabricOptimizer, ...]: ...
+
     def setup_optimizers(self, *optimizers: Optimizer) -> Union[_FabricOptimizer, tuple[_FabricOptimizer, ...]]:
         r"""Set up one or more optimizers for accelerated training.
 
@@ -409,6 +459,22 @@ def setup_optimizers(self, *optimizers: Optimizer) -> Union[_FabricOptimizer, tu
         ]
         return optimizers[0] if len(optimizers) == 1 else tuple(optimizers)
 
+    @overload
+    def setup_dataloaders(
+        self, dataloader: DataLoader, /, *, use_distributed_sampler: bool = ..., move_to_device: bool = ...
+    ) -> DataLoader: ...
+
+    @overload
+    def setup_dataloaders(
+        self,
+        first_dataloader: DataLoader,
+        second_dataloader: DataLoader,
+        /,
+        *dataloaders: DataLoader,
+        use_distributed_sampler: bool = ...,
+        move_to_device: bool = ...,
+    ) -> list[DataLoader]: ...
+
     def setup_dataloaders(
         self, *dataloaders: DataLoader, use_distributed_sampler: bool = True, move_to_device: bool = True
     ) -> Union[DataLoader, list[DataLoader]]: