Fix TypeError in contrib.reduce_on_plateau() when x64 is enabled

[stefanocortinovis]

The current implementation of reduce_on_plateau raises a TypeError when it is used with 64 bit floats enabled.

The error can be reproduced by running the following MRE:

from jax import config
config.update("jax_enable_x64", True)

import jax.numpy as jnp
from optax import contrib

updates = {"params": jnp.array(1.0)}
transform = contrib.reduce_on_plateau()
transform_state = transform.init(updates["params"])

updates, transform_state = transform.update(updates=updates, state=transform_state, loss=1.0)  # raises TypeError

This is due to the fact that, when 64 bit floats are enabled, the two branches of the jax.lax.cond at line 114 of reduce_on_plateau.py, in_cooldown and not_in_cooldown return types (int64, float32, int32) and (int32, float32, int64).

The issue can be fixed by explicitly casting the two int64 to int32.

[fabianp]

Thanks for the contribution! This looks good to me, but could you add a test so we're sure we don't break this in the future? Thanks again!

[stefanocortinovis]

I added tests for the x64 fix as parameterised test cases of the ones that were already in reduce_on_plateau_test.py. While I was at it, I've also slightly refactored the test file. Let me know if you'd like me to do something else!

[fabianp]

Amazing, thanks!