diff --git a/src/lightning/fabric/CHANGELOG.md b/src/lightning/fabric/CHANGELOG.md index de997cc4448e4..cfac76492522b 100644 --- a/src/lightning/fabric/CHANGELOG.md +++ b/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 diff --git a/src/lightning/fabric/fabric.py b/src/lightning/fabric/fabric.py index 6d76a1eb6fead..c676ad9be37c0 100644 --- a/src/lightning/fabric/fabric.py +++ b/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]]: