refactor(infer.elbo): add type hints to elbo module

[brendancooley]

Picking up on the work of @juanitorduz, this MR adds type hints to the numpyro.infer.elbo module. Along the way, I've made a few housekeeping changes to try and make it easier to maintain and extend type hints going forward.

The most important change warranting discussion is that I've introduced a ParamSpec for the parameters of a model function. ELBOs are now generics over the signature of the model/guide functions they operate on. I think that this is a nice way to (softly) enforce consistency in the *args and **kwargs passed to ELBO.loss, but there may be some subtleties that I'm missing here.

Summary of Changes

• Added comprehensive type annotations to elbo module
• Remove dormant _compute_downstream_costs code and associated tests (replaced in 28e38d8)
• Use Ruff ANN rule to require annotations in typed modules
  • exempt nested functions in numpyro.primitives with noqa
• Added numpyro._typing module for cross-module shared type aliases
• Converted some rng_key annotations to jax.Array to address typing errors and following guidance in JEP 9263
• Use shared type aliases and python 3.10+ union annotations in numpyro.handlers, with requisite import from __future__ import annotations to support python 3.9.

Previous Typing MRs (for Reference)

1. primitives
2. handlers
3. optim

[juanitorduz]

This looks great! (I approved, but we need a true core dev to approve it). I think the error

FAILED test/test_example_utils.py::test_mnist_data_load - urllib.error.HTTPError: HTTP Error 429: Too Many Requests

Is not related with these changes.

[fehiepsi]

This is awesome! I'll need to look into details and learn from your designs. :)

Please ignore the failing tests. I guess it's a github issue and will be resolved by github devs soon.

[brendancooley]

Thanks @fehiepsi and @juanitorduz for the reviews here. Still getting the hang of these typing concepts myself as you can see! Hopefully it will get easier with time and practice.

[fehiepsi]

Looks awesome, thanks @brendancooley!

[fehiepsi]

could you keep Message = MessageT in case users already use Message in their libraries?

[brendancooley]

oops yes, forgot this was pre-existing. reverted.

[fehiepsi]

does this mean that we only allow mypy to work with python 3.12?

[brendancooley]

for python 3.9, mypy complains about modern (3.10+) type annotations on our type aliases

LossT: TypeAlias = jax.Array | dict[str, jax.Array]

numpyro/infer/elbo.py:44: error: Invalid type alias: expression is not a valid type  [valid-type]
numpyro/infer/elbo.py:44: error: Unsupported left operand type for | ("type[Array]")  [operator]
Found 2 errors in 1 file (checked 88 source files)

Since python 3.9 is approaching end of life, I went ahead removed this target python_version but excluded the lint workflow from CI for the 3.9 case, which will let us use X | Y in lieu of typing.Union[X, Y] across the codebase without having to add inline ignores for the typechecker.

Alternatively, I could add an inline ignore or adopt use the legacy Union[...] type hint. Whatever you prefer!