Type Hints for Optim module and Setup MyPy as part of CI #2853

kamathhrishi · 2021-06-01T09:35:59Z

This PR adds type hints to optim module.

It addresses issue #2550

This is the first PR that addresses type hints. I will be addressing other modules one PR at a time gradually covering the codebase.

Right now its a WIP since I have to reverify some of the types and ensure all functions are covered.

kamathhrishi · 2021-06-01T13:48:17Z

pyro/optim/dct_adam.py

@@ -103,7 +104,7 @@ def step(self, closure=None):

        return loss

-    def _step_param(self, group, p):
+    def _step_param(self, group: Dict, p) -> None:


I din't get what p is and the datatype is expected?

I believe p is the parameter to be differentiated and is of type torch.Tensor (it may or may not be the narrower type torch.nn.Parameter)

kamathhrishi · 2021-06-01T13:49:22Z

pyro/optim/dct_adam.py

@@ -144,7 +145,7 @@ def _step_param(self, group, p):
            step = _transform_inverse(exp_avg / denom, time_dim, duration)
            p.data.add_(step.mul_(-step_size))

-    def _step_param_subsample(self, group, p, subsample):
+    def _step_param_subsample(self, group: Dict, p, subsample) -> None:


The expected datatype of p and subsample?

fritzo

Hi @kamathhrishi thanks so much for starting this effort! Looks great so far.

I have a few general questions for the broader effort that make sense to address in this PR:

Should we add a typechecking stage to CI? (using mypy or some other checker)
Should we specify types for internal functions like ._step()? I have only a little experience with type hints and am unsure whether they are intended more for public functions or are intended for more thorough use.
Is there an easy way to inspect undocumented code so as to determine desired type hint during this effort? E.g. can we run all unit tests and collect a log of the types passed to every function and every method? (this is more a question of your workflow during the proposed series of PRs)

kamathhrishi · 2021-06-02T01:21:25Z

Thank you

Ideally use we should. I am experimenting with mypy right now. Its great, it helped catch a lot of errors I had made during adding type hints.
Yes, we should a type checker like mypy performs a detailed analysis of literally every function to infer if the arguments and return type match.
While I am not sure, I think there is a tool for that. I will find out.

kamathhrishi · 2021-06-02T13:35:15Z

I have added mypy as part of CI. I am trying to experiment with automatically infering type hints from tests. It works with basic code with tools like pyannotate and monkeyannotate. With pyro, I unable to think how to integrate them.

fritzo · 2021-06-02T15:09:28Z

@kamathhrishi this is great, thanks for adding a mypy stage! I took a look at some of the errors. I think the first step is to configure mypy to ignore most errors (in setup.cfg?). Here are some comments on specific errors:

pyro/contrib/funsor/handlers/runtime.py:52: error: "Callable[[Type[_NT], Any, Any], _NT]" has no attribute "__defaults__"

Is this a Python version thing? I think it's fine to use a higher Python version like 3.7 or 3.8 for the mypy stage.

pyro/distributions/torch_patch.py:98: error: Need type annotation for '__all__' (hint: "__all__: List[<type>] = ...")

Should be __all__ : List[str] = ....

pyro/util.py:15: error: Skipping analyzing 'numpy': found module but no type hints or library stubs

I'm pretty sure a newer numpy supports type hints; if not there should be a stub library...

pyro/poutine/runtime.py:12: error: Need type annotation for '_PYRO_STACK' (hint: "_PYRO_STACK: List[<type>] = ...")

I believe this is _PYRO_STACK : List[Messenger].

pyro/distributions/torch_distribution.py:203: error: Definition of "log_prob" in base class "Distribution" is incompatible with definition in base class "Distribution"
pyro/distributions/torch_distribution.py:203: error: Definition of "sample" in base class "Distribution" is incompatible with definition in base class "Distribution"

We may be able to fix the Distribution methods upstream, but that will take many months before stable release.

pyro/distributions/torch_distribution.py:281: error: Need type annotation for 'arg_constraints' (hint: "arg_constraints: Dict[<type>, <type>] = ...")
pyro/distributions/torch_distribution.py:361: error: Need type annotation for 'arg_constraints' (hint: "arg_constraints: Dict[<type>, <type>] = ...")

I believe this is arg_constraints : Dict[str, torch.distributions.constraints.Constraint].

pyro/distributions/lkj.py:8: error: Module 'pyro.distributions.torch' has no attribute 'LKJCholesky'
pyro/distributions/lkj.py:8: error: Module 'pyro.distributions.torch' has no attribute 'TransformedDistribution'
...
pyro/poutine/__init__.py:4: error: Module 'pyro.poutine.handlers' has no attribute 'block'
pyro/poutine/__init__.py:4: error: Module 'pyro.poutine.handlers' has no attribute 'broadcast'

In Pyro we often use automatic import mechanisms which appear to confuse MyPy. I'm unsure the best way to resolve this in the long term, but in the short term I think it's best to ignore these errors. (Maybe in the long term we could add a script to automatically update __init__.py files to explicitly declare dynamic imports?)

pyro/poutine/collapse_messenger.py:15: error: Cannot find implementation or library stub for module named 'funsor'

Funsor is an optional dependency. I think it's safe to ignore these for now.

pyro/contrib/forecast/util.py:227: error: Name '_' already defined on line 220
...

We use _ often and I'm happy with that pattern. Let's try to configure mypy to ignore _ if possible.

.travis.yml

kamathhrishi · 2021-06-02T15:27:06Z

@fritzo I accidentally ran it for the whole module. Right now I have added the typechecks only for the optim module. I have had some errors in a few places which I ignored by using #type:ignore keyword. I will let you know when I find a fix for automatic imports.

kamathhrishi · 2021-06-02T15:28:10Z

Also you would like me to skip private functions with a _ from mypy or to completely not use type annotations. I think it would help any new developer to read the codebase easier if added to private functions.

fritzo · 2021-06-02T15:49:12Z

I agree with your argument that it would help developers to include private functions in type checking.

pyro/optim/adagrad_rmsprop.py

kamathhrishi · 2021-06-03T12:33:29Z

I agree with your argument that it would help developers to include private functions in type checking.

So I will be adding the type hints for them.

.mypy.ini

kamathhrishi · 2021-06-03T14:00:01Z

@fritzo There are some functions I haven't annotated yet mainly dct_adam. I will do that. Is there anything else to do done before it is ready to be merged?
@fehiepsi Let me know if you see some other changes to be made.

fritzo

LGTM, this is a great first PR!

@kamathhrishi is it possible to enable type checking incrementally for new files? E.g. in #2856 I have added type hints. If I merge that PR after this one, how would I update setup.cfg to enable checking for the pyro/ops/streaming.py and tests/ops/test_streaming.py?

kamathhrishi · 2021-06-03T16:33:15Z

@fritzo Thank you, looking forward to adding it to the rest of the modules.
So right now I have disabled type checking for the entire pyro.ops because this PR had added type hints only for optim module.
What I would suggest you do is remove the rule set for pyro.ops to ignore errors. Then individually ignore type errors from files you haven't touched. Now I ran mypy on ops folder and noticed there were mypy errors for einsum.py, contract.py and tensor_utils.py so I would add the rule

[mypy-pyro.ops.contract]
ignore_errors = True

and similarly for einsum.py, contract.py.

This is one solution, a better way to do is to somehow ignore errors completely for ops except for streaming.py. I tried something like

[mypy-pyro.ops.*]
ignore_errors = False 

[mypy-pyro.ops.einsum]
ignore_errors = True`

Din't work as expected.

Further, I don't think you should perform type checking on tests. Just modules.

fritzo · 2021-06-03T17:00:06Z

Thanks @kamathhrishi, I've updated #2856 as you suggested. The only hitch was that I needed to move all the warn_incomplete_stub = True into the main [mypy] section to avoid errors.

kamathhrishi added 2 commits June 1, 2021 14:20

Add type hints to optim module

64d0481

Changes according to linter

bd77ae8

kamathhrishi commented Jun 1, 2021

View reviewed changes

Completed and correctedd some type hints

b27dc4f

kamathhrishi changed the title ~~[WIP] Type Hints for Optim module~~ Type Hints for Optim module Jun 1, 2021

fritzo added the usability label Jun 1, 2021

fritzo reviewed Jun 1, 2021

View reviewed changes

kamathhrishi added 5 commits June 2, 2021 16:25

Changes in accordance to mypy

a3aae4f

Make lint happy

b39b7a8

Remove mypy cache

7823aee

Add mypy cache to gitignore

8c8839a

Update .travis.yml

e06ecc6

fritzo reviewed Jun 2, 2021

View reviewed changes

.travis.yml Outdated Show resolved Hide resolved

Update .travis.yml

5964d78

Update .travis.yml

9d08fb4

fehiepsi reviewed Jun 2, 2021

View reviewed changes

pyro/optim/adagrad_rmsprop.py Outdated Show resolved Hide resolved

kamathhrishi added 7 commits June 3, 2021 06:43

Add mypy config

6d58ea4

Update .travis.yml

6cf93ef

Fix linter errors

e57d897

Update .travis.yml

7b3f9f2

Update .travis.yml

8f69ec8

Update .mypy.ini

c33cf1d

Update __init__.py

3a61edc

kamathhrishi changed the title ~~Type Hints for Optim module~~ Type Hints for Optim module and Setup MyPy as part of CI Jun 3, 2021

Update __init__.py

0fd48b5

kamathhrishi added 2 commits June 3, 2021 18:22

Remove optional when default argument isnt None

e36197a

Update .mypy.ini

f074a85

fritzo reviewed Jun 3, 2021

View reviewed changes

.mypy.ini Outdated Show resolved Hide resolved

kamathhrishi added 4 commits June 3, 2021 19:27

Update .mypy.ini

86224aa

Delete .mypy.ini

64cae9e

Update setup.cfg

21f470f

Update setup.cfg

d19d8a7

fritzo approved these changes Jun 3, 2021

View reviewed changes

fritzo merged commit 145141c into pyro-ppl:dev Jun 3, 2021

fritzo mentioned this pull request Jun 7, 2021

Using Python type hints #2550

Open

23 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Type Hints for Optim module and Setup MyPy as part of CI #2853

Type Hints for Optim module and Setup MyPy as part of CI #2853

kamathhrishi commented Jun 1, 2021 •

edited

Loading

kamathhrishi Jun 1, 2021

fritzo Jun 1, 2021

kamathhrishi Jun 1, 2021

fritzo left a comment

kamathhrishi commented Jun 2, 2021 •

edited by fritzo

Loading

kamathhrishi commented Jun 2, 2021 •

edited by fritzo

Loading

fritzo commented Jun 2, 2021

kamathhrishi commented Jun 2, 2021 •

edited

Loading

kamathhrishi commented Jun 2, 2021 •

edited

Loading

fritzo commented Jun 2, 2021

kamathhrishi commented Jun 3, 2021

kamathhrishi commented Jun 3, 2021 •

edited

Loading

fritzo left a comment

kamathhrishi commented Jun 3, 2021

fritzo commented Jun 3, 2021

Type Hints for Optim module and Setup MyPy as part of CI #2853

Type Hints for Optim module and Setup MyPy as part of CI #2853

Conversation

kamathhrishi commented Jun 1, 2021 • edited Loading

kamathhrishi Jun 1, 2021

Choose a reason for hiding this comment

fritzo Jun 1, 2021

Choose a reason for hiding this comment

kamathhrishi Jun 1, 2021

Choose a reason for hiding this comment

fritzo left a comment

Choose a reason for hiding this comment

kamathhrishi commented Jun 2, 2021 • edited by fritzo Loading

kamathhrishi commented Jun 2, 2021 • edited by fritzo Loading

fritzo commented Jun 2, 2021

kamathhrishi commented Jun 2, 2021 • edited Loading

kamathhrishi commented Jun 2, 2021 • edited Loading

fritzo commented Jun 2, 2021

kamathhrishi commented Jun 3, 2021

kamathhrishi commented Jun 3, 2021 • edited Loading

fritzo left a comment

Choose a reason for hiding this comment

kamathhrishi commented Jun 3, 2021

fritzo commented Jun 3, 2021

kamathhrishi commented Jun 1, 2021 •

edited

Loading

kamathhrishi commented Jun 2, 2021 •

edited by fritzo

Loading

kamathhrishi commented Jun 2, 2021 •

edited by fritzo

Loading

kamathhrishi commented Jun 2, 2021 •

edited

Loading

kamathhrishi commented Jun 2, 2021 •

edited

Loading

kamathhrishi commented Jun 3, 2021 •

edited

Loading