Use type hints whenever possible

[neerajprad]

Since we are using python 3, we can start taking advantage of type hints wherever possible.

[jpchen]

we should probably do this in pyro too

[akihironitta]

@neerajprad @jpchen Could I work on this issue?

[neerajprad]

@akihironitta - The codebase and number of contributors to numpyro core is small enough that this is low priority for us at the moment. When I created the issue, I was thinking that we could do this in a piecemeal fashion as we were refactoring the codebase. If you are interested, my suggestion would be to start with one of the other tasks and refactor to use type hints on the internal code that your PR touches.

[juanitorduz]

@fehiepsi is this something you would like to add? I have done it for other libraries and we could tackle this on PR per module at a time (as in pyro-ppl/pyro#2550) . We could also use https://github.com/patrick-kidger/jaxtyping (which I use and like quite a lot ;) )

[fehiepsi]

I think it would be nice to have. Because this would require a significant amount of work, please make the PRs modular and incremental.

[juanitorduz]

I think it would be nice to have. Because this would require a significant amount of work, please make the PRs modular and incremental.

Ok! We can test a simple first PR for a module and see if we like it 👍