Add option to create_optimizer for colocate_gradients_with_ops
Add option to create_optimizer for colocate_gradients_with_ops

Set this to True if you are using multiple cores per replica. See "num_cores_per_replica": https://www.tensorflow.org/api_docs/python/tf/compat/v1/estimator/tpu/TPUConfig

PiperOrigin-RevId: 298991782

google-research/albert,google-research/albert