diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index a3655699669..e39631d634f 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -121,9 +121,17 @@ cuda_py_test(
     name = "device_util_test",
     srcs = ["device_util_test.py"],
     deps = [
+        ":combinations",
         ":device_util",
+        ":multi_worker_test_base",
+        ":multi_worker_util",
+        "//tensorflow/core:protos_all_py",
         "//tensorflow/python:client_testlib",
+        "//tensorflow/python:extra_py_tests_deps",
         "//tensorflow/python:framework_ops",
+        "//tensorflow/python:training_server_lib",
+        "//tensorflow/python/eager:context",
+        "@absl_py//absl/testing:parameterized",
     ],
 )
 
diff --git a/tensorflow/python/distribute/device_util.py b/tensorflow/python/distribute/device_util.py
index 7f32ed39aed..0aee34e33ae 100644
--- a/tensorflow/python/distribute/device_util.py
+++ b/tensorflow/python/distribute/device_util.py
@@ -19,6 +19,7 @@ from __future__ import division
 from __future__ import print_function
 
 from tensorflow.python.eager import context
+from tensorflow.python.framework import config
 from tensorflow.python.framework import device as tf_device
 from tensorflow.python.framework import ops
 
@@ -55,8 +56,16 @@ def canonicalize(d, default=None):
   result = tf_device.DeviceSpec(
       replica=0, task=0, device_type="CPU", device_index=0)
   if ops.executing_eagerly_outside_functions():
-    # The default job is localhost if eager execution is enabled
-    result = result.replace(job="localhost")
+    # Try to deduce job, replica and task in case it's in a multi worker setup.
+    # TODO(b/151452748): Using list_logical_devices is not always safe since it
+    # may return remote devices as well, but we're already doing this elsewhere.
+    host_cpu = tf_device.DeviceSpec.from_string(
+        config.list_logical_devices("CPU")[0].name)
+    if host_cpu.job:
+      result = result.make_merged_spec(host_cpu)
+    else:
+      # The default job is localhost if eager execution is enabled
+      result = result.replace(job="localhost")
   if default:
     # Overrides any defaults with values from the default device if given.
     result = result.make_merged_spec(
diff --git a/tensorflow/python/distribute/device_util_test.py b/tensorflow/python/distribute/device_util_test.py
index 2f0d7ed3b31..df53fe0288a 100644
--- a/tensorflow/python/distribute/device_util_test.py
+++ b/tensorflow/python/distribute/device_util_test.py
@@ -18,16 +18,27 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+from absl.testing import parameterized
+
+from tensorflow.core.protobuf import tensorflow_server_pb2
+from tensorflow.python.distribute import combinations
 from tensorflow.python.distribute import device_util
+from tensorflow.python.distribute import multi_worker_test_base
 from tensorflow.python.eager import context
 from tensorflow.python.framework import ops
-from tensorflow.python.framework import test_util
 from tensorflow.python.platform import test
+from tensorflow.python.training import server_lib
 
 
-class DeviceUtilTest(test.TestCase):
+class DeviceUtilTest(test.TestCase, parameterized.TestCase):
 
-  @test_util.run_deprecated_v1
+  def setUp(self):
+    super(DeviceUtilTest, self).setUp()
+    context._reset_context()  # pylint: disable=protected-access
+
+  @combinations.generate(
+      combinations.combine(mode="graph")
+  )
   def testCurrentDeviceWithGlobalGraph(self):
     with ops.device("/cpu:0"):
       self.assertEqual(device_util.current(), "/device:CPU:0")
@@ -51,11 +62,16 @@ class DeviceUtilTest(test.TestCase):
         self.assertEqual(device_util.current(),
                          "/job:localhost/replica:0/task:0/device:CPU:0")
 
-  @test_util.run_deprecated_v1
-  def testCanonicalizeWithoutDefaultDevice(self):
-    self.assertEqual(
-        device_util.canonicalize("/cpu:0"),
-        "/replica:0/task:0/device:CPU:0")
+  @combinations.generate(combinations.combine(mode=["graph", "eager"]))
+  def testCanonicalizeWithoutDefaultDevice(self, mode):
+    if mode == "graph":
+      self.assertEqual(
+          device_util.canonicalize("/cpu:0"),
+          "/replica:0/task:0/device:CPU:0")
+    else:
+      self.assertEqual(
+          device_util.canonicalize("/cpu:0"),
+          "/job:localhost/replica:0/task:0/device:CPU:0")
     self.assertEqual(
         device_util.canonicalize("/job:worker/cpu:0"),
         "/job:worker/replica:0/task:0/device:CPU:0")
@@ -63,6 +79,22 @@ class DeviceUtilTest(test.TestCase):
         device_util.canonicalize("/job:worker/task:1/cpu:0"),
         "/job:worker/replica:0/task:1/device:CPU:0")
 
+  @combinations.generate(combinations.combine(mode=["eager"]))
+  def testCanonicalizeWithoutDefaultDeviceCollectiveEnabled(self):
+    cluster_spec = server_lib.ClusterSpec(
+        multi_worker_test_base.create_cluster_spec(
+            has_chief=False, num_workers=1, num_ps=0, has_eval=False))
+    server_def = tensorflow_server_pb2.ServerDef(
+        cluster=cluster_spec.as_cluster_def(),
+        job_name="worker",
+        task_index=0,
+        protocol="grpc",
+        port=0)
+    context.context().enable_collective_ops(server_def)
+    self.assertEqual(
+        device_util.canonicalize("/cpu:0"),
+        "/job:worker/replica:0/task:0/device:CPU:0")
+
   def testCanonicalizeWithDefaultDevice(self):
     self.assertEqual(
         device_util.canonicalize("/job:worker/task:1/cpu:0", default="/gpu:0"),