Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix typing of poutine.trace(model).get_trace() #3334

Merged
merged 2 commits into from
Feb 27, 2024
+14 −14
Merged

Fix typing of poutine.trace(model).get_trace() #3334

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e0e3a82
Fix typing of poutin.trace(fn)
fritzo Feb 27, 2024
401e55b
Refine type hint for ReparamMessenger
fritzo Feb 27, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions pyro/poutine/handlers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@
from pyro.poutine.scale_messenger import ScaleMessenger
from pyro.poutine.seed_messenger import SeedMessenger
from pyro.poutine.substitute_messenger import SubstituteMessenger
from pyro.poutine.trace_messenger import TraceMessenger
from pyro.poutine.trace_messenger import TraceHandler, TraceMessenger
from pyro.poutine.uncondition_messenger import UnconditionMessenger

if TYPE_CHECKING:
Expand Down Expand Up @@ -467,26 +467,26 @@ def substitute( # type: ignore[empty-body]

@overload
def trace(
fn: None = ...,
fn: None = None,
graph_type: Optional[Literal["flat", "dense"]] = None,
param_only: Optional[bool] = None,
) -> TraceMessenger: ...


@overload
def trace(
fn: Callable[_P, _T] = ...,
fn: Callable[_P, _T],
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I had to remove the default to avoid the mypy error:

pyro/poutine/handlers.py:469: error: Overloaded function signatures 1 and 2 overlap with incompatible return types  [overload-overlap]

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Interesting, seems like some weird behavior from mypy

graph_type: Optional[Literal["flat", "dense"]] = None,
param_only: Optional[bool] = None,
) -> Callable[_P, _T]: ...
) -> TraceHandler[_P, _T]: ...


@_make_handler(TraceMessenger)
def trace( # type: ignore[empty-body]
fn: Optional[Callable[_P, _T]] = None,
graph_type: Optional[Literal["flat", "dense"]] = None,
param_only: Optional[bool] = None,
) -> Union[TraceMessenger, Callable[_P, _T]]: ...
) -> Union[TraceMessenger, TraceHandler[_P, _T]]: ...


@overload
Expand Down
Loading