Skip to content

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

Open
dai123x wants to merge 2 commits into
Lightning-AI:masterfrom
dai123x:fabric-setup-overloads
Open

dai123x wants to merge 2 commits into
Lightning-AI:masterfrom
dai123x:fabric-setup-overloads

Conversation

@dai123x

@dai123x dai123x commented Oct 7, 2026

Copy link
Copy Markdown

What does this PR do?

Adds typing.overload signatures for Fabric.setup, Fabric.setup_optimizers and Fabric.setup_dataloaders, so the declared return type follows the number of arguments.

Fixes #21988

Fabric.setup previously returned Any, and the other two methods returned a union of the single- and the multi-object case, so type checkers lost the shape of the most common calls and every access on the result needed a cast or an explicit annotation (examples in #21988). Annotations only — no runtime behavior change.

call before after
fabric.setup(model) Any _FabricModule
fabric.setup(model, optimizer) Any tuple[_FabricModule, _FabricOptimizer]
fabric.setup(model, optimizer, scheduler=lr) Any tuple[_FabricModule, _FabricOptimizer, LRScheduler]
fabric.setup(model, opt1, opt2) Any tuple[_FabricModule | _FabricOptimizer | LRScheduler, ...]
fabric.setup_optimizers(opt) union _FabricOptimizer
fabric.setup_optimizers(opt1, opt2) union tuple[_FabricOptimizer, ...]
fabric.setup_dataloaders(loader) union DataLoader
fabric.setup_dataloaders(loader1, loader2) union list[DataLoader]

Design notes (kept as close to the sketch in #21988 as possible):

  • The "multiple objects" overloads require at least two positional optimizers/dataloaders instead of a plain *args. Calls the runtime rejects (e.g. setup_optimizers() with no arguments) then fall through to the implementation signature instead of being typed with a shape the runtime never returns.
  • setup(module=model) keeps working: module stays keyword-callable, while the optimizer arguments are positional-only, matching the runtime signature.
  • The multi-optimizer union also covers the scheduler, which the runtime appends last when one is provided.

For reviewers

  • Checked with mypy 1.20.0 (the version pinned in requirements/typing.txt): no new errors. A consumer snippet resolves every shape in the table above, verified with reveal_type.
  • pytest tests/tests_fabric/test_fabric.py tests/tests_fabric/test_wrappers.py passes locally (Windows, CPU torch 2.12.0): 112 passed, 66 skipped (distributed/TPU skips), 0 failures.
  • No runtime tests added since behavior is unchanged; no docs or CHANGELOG update needed for an annotation-only change.
Before submitting
  • Was this discussed/agreed via a GitHub issue? Yes — Add overloads to Fabric.setup, setup_optimizers, and setup_dataloaders for precise return types #21988 describes the problem and sketches these overloads.
  • Did you read the contributor guideline, Pull Request section?
  • Did you make sure your PR does only one thing, instead of bundling different changes together?
  • Did you make sure to update the documentation with your changes? (if necessary) — not necessary, annotations only
  • Did you write any new necessary tests? (not for typos and docs) — no new runtime tests; behavior unchanged, typing verified with mypy as described above
  • Did you verify new and existing tests pass locally with your changes?
  • Did you list all the breaking changes introduced by this pull request? — none

No breaking changes — annotations only.

…izers and setup_dataloaders

Fabric.setup returned Any and the other two methods returned a union of the
single- and the multi-object case, so type checkers lost the shape of the most
common calls (Lightning-AI#21988). Add typing.overload signatures that follow the number of
arguments. Annotations only; runtime behavior is unchanged.

Signed-off-by: dai123x <daixuan26@outlook.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add overloads to Fabric.setup, setup_optimizers, and setup_dataloaders for precise return types

1 participant