From f8c912d2b70280752563f6bbbf626c5e5ea72b6a Mon Sep 17 00:00:00 2001 From: Zhenyu Tan Date: Mon, 22 Jul 2019 14:14:12 -0700 Subject: [PATCH] Add colocation back for optimizer v2. PiperOrigin-RevId: 259403870 --- tensorflow/python/keras/optimizer_v2/optimizer_v2.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py index 039e2b4cea7..f053d856dd3 100644 --- a/tensorflow/python/keras/optimizer_v2/optimizer_v2.py +++ b/tensorflow/python/keras/optimizer_v2/optimizer_v2.py @@ -474,9 +474,12 @@ class OptimizerV2(trackable.Trackable): update_ops = [] with backend.name_scope(name or self._name): for grad, var in grads_and_vars: - scope_name = ("" if ops.executing_eagerly_outside_functions() else - "_" + var.op.name) - with backend.name_scope("update" + scope_name): + scope_name = ("update" if ops.executing_eagerly_outside_functions() else + "update_" + var.op.name) + # Colocate the update with variables to avoid unnecessary communication + # delays. See b/136304694. + with backend.name_scope( + scope_name), distribution.extended.colocate_vars_with(var): update_ops.extend( distribution.extended.update( var, apply_grad_to_update_var, args=(grad,), group=False))