Add option for mixed precision training

[dustinvtran]

Deterministic and BatchEnsemble baselines cast data as bfloat16 by default with TPUs. Following the cloud TPU imagenet example, we need to set a policy if we want to maintain bfloat16 for the activations, etc.

  if _USE_BFLOAT16:
    policy = tf.keras.mixed_precision.experimental.Policy('mixed_bfloat16')
    tf.keras.mixed_precision.experimental.set_policy(policy)

It should be easy enough to set up a boolean flag to use bfloat16 with that policy or otherwise operate in float32.