Pass value arg to optax, allowing use of reduce_on_plateau

[zmbc]

Addresses #1955 using the approach from this comment: #1955 (comment)

[zmbc]

@fehiepsi Let me know what you think! I haven't tested this code, and I might need some help with adding types as well.

[fehiepsi]

You can introduce an attribute update_with_value (false by default) to control the behavior here

if self.update_with_value=True:
    opt_state = self.update_fn(i, g, opt_state, value)
else:
    opt_state = self.update_fn(i, g, opt_state)

you can then use the typing

self.update_fn: Union[Callable[[ArrayLike, _Params, _OptState], _OptState], Callable[[ArrayLike, _Params, _OptState, ArrayLike], _OptState]]

In optax optimizer, you can set it to True

numpyro_optim = _NumPyroOptim(lambda x, y, z: (x, y, z), init_fn, update_fn, get_params_fn)
numpyro_optim.update_with_value = True
return numpyro_optim

[zmbc]

@fehiepsi I can't easily get this typing to pass mypy. I'd need to use some kind of custom TypeGuard to narrow this Union type down to the specific Callable signature at runtime before it is called. Do you want me to do that? Or would you rather I changed this to simply Callable[..., _OptState]?

[zmbc]

@fehiepsi thanks for your comments, want to take another look?

[fehiepsi]

Could you add a simple test for this?