Fix accidental type promotion in adam and adamax optimizers.
Fix accidental type promotion in adam and adamax optimizers.

google/jax,tensorflow/probability,google/jax,google/jax,tensorflow/probability,google/jax