MyPy Errors in Distributions Module

[juanitorduz]

Bug Description

In the PR, adding type hints to the distributions module https://github.com/pyro-ppl/numpyro/pull/2032, the CI did not check for consistency, and we have many lint errors.

Steps to Reproduce

Steps to reproduce the behavior.

We must always add the module into the mypy config in the pyproject.toml file as

[[tool.mypy.overrides]]
module = [
    "numpyro.contrib.control_flow.*",       # types missing
    "numpyro.contrib.funsor.*",             # types missing
    "numpyro.contrib.hsgp.*",
    "numpyro.contrib.stochastic_support.*",
    "numpyro.diagnostics.*",
    "numpyro.distributions.*",    # <-- This line was missing!
    "numpyro.handlers.*",
    "numpyro.infer.elbo.*",
    "numpyro.optim.*",
    "numpyro.primitives.*",
    "numpyro.patch.*",
    "numpyro.util.*",
]
ignore_errors = false

When I add this line and do the mypy checks, we have many errors

...
Found 1043 errors in 14 files (checked 88 source files)

Expected Behavior

We should add the line "numpyro.distributions.*" to the pyproject.toml file as above and fix all the errors. As there are many, we can do it in small PRs.

FYI: @Qazalbash