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 Any pointers here?

[fehiepsi]

Sorry I missed the message. I think you can define a protocol for the purpose

class _UpdateFn(Protocol):
    def __call__(self, arr: ArrayLike, params: _Params, opt_state: _OptState, *args: ArrayLike) -> _OptState:
        ...

[tillahoffmann]

It might be possible to reuse the optax protocol here, although the OptStates may not be compatible.

[zmbc]

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

[fehiepsi]

Could you add a simple test for this?

[fehiepsi]

@zmbc, to test the optimizer, you can simply add the new one to this line

numpyro/test/test_optimizers.py

Line 29 in d393f5a

(optax.sgd, (1e-2,), {}),

[zmbc]

@fehiepsi I think I've got this all working except that the test test_numpyrooptim_no_double_jit is failing on this optimizer with

>       assert my_fn_calls == 1
E       assert 2 == 1

test/test_optimizers.py:125: AssertionError

My understanding of JAX and JIT is pretty basic, any ideas why this might be happening?

[fehiepsi]

Thanks for supporting those new optimizers, @zmbc!

[fehiepsi]

test_numpyrooptim_no_double_jit is failing

I think the init state and the later states of the new optimizer have different dtype. Could you double check? If it is the case, you can assert <= 2 for that specific case.

[zmbc]

@fehiepsi You are correct and I'm glad I asked; that would have taken me a long time to figure out!

The difference in type is small (see comment):

(
    Traced<ShapedArray(int32[], weak_type=True)> with <DynamicJaxprTrace>,
    (
        Traced<ShapedArray(float32[10])> with <DynamicJaxprTrace>,
        (
            (
                ScaleByAdamState(
                    count=Traced<ShapedArray(int32[])> with <DynamicJaxprTrace>,
                    mu=Traced<ShapedArray(float32[10])> with <DynamicJaxprTrace>,
                    nu=Traced<ShapedArray(float32[10])> with <DynamicJaxprTrace>
                ),
                EmptyState()
            ),
            ReduceLROnPlateauState(
                scale=Traced<ShapedArray(float32[])> with <DynamicJaxprTrace>,
                best_value=Traced<ShapedArray(float32[], weak_type=True)> with <DynamicJaxprTrace>,
                plateau_count=Traced<ShapedArray(int32[])> with <DynamicJaxprTrace>,
                cooldown_count=Traced<ShapedArray(int32[])> with <DynamicJaxprTrace>,
                count=Traced<ShapedArray(int32[])> with <DynamicJaxprTrace>,
                avg_value=Traced<ShapedArray(float32[],
                  weak_type=True, # This is only present on the first call
                )> with <DynamicJaxprTrace>
            )
        )
    )
)

Does that seem correct to you? I'll update the test.

[fehiepsi]

Could you add value=loss_val here

numpyro/numpyro/contrib/einstein/steinvi.py

Line 470 in 374cd89

optim_state = self.optim.update(grads, optim_state)

to make tests pass?

[fehiepsi]

There is a conflict. Could you merge with the master branch?

[zmbc]

Sorry about that, should be resolved now!