Fix error with x64 loss

[stefanocortinovis]

This PR fixes a bug for contrib.reduce_on_plateau() when 64-bit floats are enabled:

• The current implementation of contrib.reduce_on_plateau() results in a TypeError: true_fun and false_fun output must have identical types due to a jax.lax.cond using functions with different return types when the value passed to transform.update() is a 64-bit scalar tensor.
• The inconsistency in the types is due to this average value computation.
• The bug was not picked up by current tests because, in the x64 tests, values were not explicitly cast to 64-bit tensors.
• The new implementation corrects the issue with the tests and solves the bug by explicitly casting value passed to transform.update() to 32-bits where needed.

[fabianp]

@stefanocortinovis please ping me once the changes have been implemented

[stefanocortinovis]

@fabianp done! I will just need to add a couple of small changes when #975 gets merged.