Add types optim module

[juanitorduz]

Add types optim module

[fehiepsi]

why use Any for _Params but not for OptState

[juanitorduz]

If I leave _Params = TypeVar("_Params") I get a MyPy error;

numpyro/optim.py:103: error: Type variable "numpyro.optim._Params" is unbound

This line is

https://github.com/juanitorduz/numpyro/blob/cd364d84308f253f0d0a251bdb0f36daaa38cfd4/numpyro/optim.py#L103-L107

    def eval_and_update(
        self,
        fn: Callable[[Any], tuple],
        state: _IterOptState,
        forward_mode_differentiation: bool = False,
    ) -> tuple[tuple[Any, Any], _IterOptState]:
        ...
        params: _Params = self.get_params(state).   # <- HERE!
        (out, aux), grads = _value_and_grad(
            fn, x=params, forward_mode_differentiation=forward_mode_differentiation
        )

The problem is that TypeVars are fine as long as they are part of the input and output. However, in our case, the variable params is created inside the functions and this is something MyPy does not like. It requires either a bound (a subset of possible types) or a concrete type.

This is one of the references that helped me clarify this issue https://stackoverflow.com/questions/68603585/mypy-why-does-typevar-not-work-without-bound-specified

[fehiepsi]

How about just using `params = self.get_params(state)?

[juanitorduz]

Unfortunately this does not work (it was my first try) :(

numpyro/optim.py:103: error: Need type annotation for "params"  [var-annotated]
numpyro/optim.py:125: error: Type variable "numpyro.optim._Params" is unbound  [valid-type]

The reason is because of the logic above. We can not use TypeVar internally :/

[juanitorduz]

I have made another attempt which I think serves our purpose in de835c6 I had to skip one type check in

    def get_params(opt_state: _MinimizeState) -> _Params:  # type: ignore[type-var]. # <- HERE
        flat_params, unravel_fn = opt_state
        return unravel_fn(flat_params)

Better?

[fehiepsi]

I guess we can use Any for both to avoid confusion.

[juanitorduz]

Sure! Changed in e24a7a2 :)

[fehiepsi]

LGTM pending the above comment. Thanks, @juanitorduz!

[juanitorduz]

I had to specify a bound or a specific type and this is currently the "best" solution for pytrees, see jax-ml/jax#3340 (comment)

[juanitorduz]

If I leave _Params = TypeVar("_Params") I get a MyPy error;

numpyro/optim.py:103: error: Type variable "numpyro.optim._Params" is unbound

This line is

https://github.com/juanitorduz/numpyro/blob/cd364d84308f253f0d0a251bdb0f36daaa38cfd4/numpyro/optim.py#L103-L107

    def eval_and_update(
        self,
        fn: Callable[[Any], tuple],
        state: _IterOptState,
        forward_mode_differentiation: bool = False,
    ) -> tuple[tuple[Any, Any], _IterOptState]:
        ...
        params: _Params = self.get_params(state).   # <- HERE!
        (out, aux), grads = _value_and_grad(
            fn, x=params, forward_mode_differentiation=forward_mode_differentiation
        )

The problem is that TypeVars are fine as long as they are part of the input and output. However, in our case, the variable params is created inside the functions and this is something MyPy does not like. It requires either a bound (a subset of possible types) or a concrete type.

This is one of the references that helped me clarify this issue https://stackoverflow.com/questions/68603585/mypy-why-does-typevar-not-work-without-bound-specified

[juanitorduz]

Ups! I had replied the comment but forgotten to hit the "reply button" 🙈