mixed precision support on GPUs

[LiweiPeng]

Hi, To speed up training on V100 GPUs, I'd like to run mesh tf using mixed precision. While TensorFlow has an easy to use automatic mixed precision feature, it requires the optimizer to be a tf.train.Optimizer. This won't work on mesh tf's optimizers.

My question is: how can I use mixed precision on GPUs with mesh tf? If not supported yet, can you add some support for this? Thanks.