Ensure optimizers return updates of same dtype as params.

Ensure optimizers return updates of same dtype as params. #4591