1 files changed, 3 insertions, 16 deletions
diff --git a/tensorflow/core/common_runtime/device_mgr.cc b/tensorflow/core/common_runtime/device_mgr.cc
index 471463fc8b..0a4e0afc87 100644
--- a/tensorflow/core/common_runtime/device_mgr.cc
+++ b/tensorflow/core/common_runtime/device_mgr.cc
@@ -30,22 +30,9 @@ DeviceMgr::DeviceMgr(const std::vector<Device*>& devices)
     devices_.push_back(d);
 
     // Register under the (1) full name, (2) canonical name, and (3) local name.
-    string full_name = d->name();
-    device_map_[CopyToBackingStore(full_name)] = d;
-
-    // TODO(b/62909072): Upgrade device_map_ to a better data structure.
-    DeviceNameUtils::ParsedName parsed_name = d->parsed_name();
-    if (parsed_name.has_job && parsed_name.has_replica &&
-        parsed_name.has_task && parsed_name.has_type && parsed_name.has_id) {
-      string canonical_name = DeviceNameUtils::FullName(
-          parsed_name.job, parsed_name.replica, parsed_name.task,
-          parsed_name.type, parsed_name.id);
-      device_map_[CopyToBackingStore(canonical_name)] = d;
-
-      string legacy_name = DeviceNameUtils::LegacyName(
-          parsed_name.job, parsed_name.replica, parsed_name.task,
-          parsed_name.type, parsed_name.id);
-      device_map_[CopyToBackingStore(legacy_name)] = d;
+    for (const string& name :
+         DeviceNameUtils::GetNamesForDeviceMappings(d->parsed_name())) {
+      device_map_[CopyToBackingStore(name)] = d;
     }
     string lname = DeviceNameUtils::LocalName(d->name());
     device_map_[CopyToBackingStore(lname)] = d;