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

[fritzo]

Addresses #3332 #2550

This proposes to replace handler types with more explicit handlers:

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

so mypy can understand code like poutine.trace(model).get_trace().

If @ordabayevy approves of this change, I can similarly change all handlers in this PR.

[fritzo]

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]

[ordabayevy]

Interesting, seems like some weird behavior from mypy

[ordabayevy]

This makes total sense. Should we also do this for the ReparamHandler even though it doesn't have any methods defined?

[ordabayevy]

Interesting, seems like some weird behavior from mypy

[fritzo]

Should we also do this for the ReparamHandler...?

Yes, I think we should be as precise as is practical. Beyond mypy correctness checking, the propagated type hints are used by vs code to show documentation on hover. So if we can narrow the type hint to a specific class, vs code will show the documentation of that variable. I've been using hover-docs more often than sphinx docs lately.

I'll go ahead and make the remaining changes, then re-submit for review. Thanks for taking a look!

[ordabayevy]

LGTM. Thank you!