Partially type annotate distributions

[ordabayevy]

No description provided.

[ordabayevy]

I can't reproduce the errors from the CI run on my machine (installed Python 3.8, upgraded torch and mypy):

mypy --install-types --non-interactive pyro scripts tests
pyro/distributions/torch_distribution.py:315: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
pyro/distributions/torch_distribution.py:400: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
pyro/distributions/rejector.py:[22](https://github.com/pyro-ppl/pyro/actions/runs/9033224237/job/24823084113?pr=3367#step:5:23): error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
pyro/distributions/polya_gamma.py:37: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
Found 14 errors in 8 files (checked 550 source files)
pyro/distributions/nanmasked.py:27: error: Call to abstract method "log_prob" of "Distribution" with trivial body via super() is unsafe  [safe-super]
pyro/distributions/nanmasked.py:27: note: Error code "safe-super" not covered by "type: ignore" comment
pyro/distributions/nanmasked.py:68: error: Call to abstract method "log_prob" of "Distribution" with trivial body via super() is unsafe  [safe-super]
pyro/distributions/nanmasked.py:68: note: Error code "safe-super" not covered by "type: ignore" comment
pyro/distributions/mixture.py:64: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
pyro/distributions/improper_uniform.py:44: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
pyro/distributions/empirical.py:49: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
pyro/distributions/hmm.py:495: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
pyro/distributions/hmm.py:814: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
pyro/distributions/hmm.py:1007: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
pyro/distributions/hmm.py:1156: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
pyro/distributions/hmm.py:1[24](https://github.com/pyro-ppl/pyro/actions/runs/9033224237/job/24823084113?pr=3367#step:5:25)1: error: Need type annotation for "arg_constraints" (hint: "arg_constraints: Dict[<type>, <type>] = ...")  [var-annotated]
make: *** [Makefile:24: lint] Error 1