diff --git a/tensorflow/core/framework/dataset.cc b/tensorflow/core/framework/dataset.cc
index c851af9a5c4..7e7bfafea80 100644
--- a/tensorflow/core/framework/dataset.cc
+++ b/tensorflow/core/framework/dataset.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/resource.h"
+#include "tensorflow/core/platform/strcat.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 
 namespace tensorflow {
@@ -527,6 +528,26 @@ DatasetBaseIterator::DatasetBaseIterator(const BaseParams& params)
     : params_(params) {
   params_.dataset->Ref();
   VLOG(2) << prefix() << " constructor";
+  traceme_metadata_ = strings::StrCat("id=", id_);
+  if (parent_) {
+    strings::StrAppend(&traceme_metadata_, ",parent_id=", parent_id_);
+  }
+  strings::StrAppend(&traceme_metadata_, ",shapes=");
+  auto& shapes = output_shapes();
+  for (int i = 0; i < shapes.size(); ++i) {
+    if (i > 0) {
+      strings::StrAppend(&traceme_metadata_, " ");
+    }
+    strings::StrAppend(&traceme_metadata_, shapes.at(i).DebugString());
+  }
+  strings::StrAppend(&traceme_metadata_, ",types=");
+  auto& types = output_dtypes();
+  for (int i = 0; i < types.size(); ++i) {
+    if (i > 0) {
+      strings::StrAppend(&traceme_metadata_, " ");
+    }
+    strings::StrAppend(&traceme_metadata_, DataTypeString(types.at(i)));
+  }
 }
 
 DatasetBaseIterator::~DatasetBaseIterator() {
@@ -535,11 +556,7 @@ DatasetBaseIterator::~DatasetBaseIterator() {
 }
 
 string DatasetBaseIterator::BuildTraceMeName() {
-  string result = strings::StrCat(params_.prefix, "#id=", id_);
-  if (parent_) {
-    strings::StrAppend(&result, ",parent_id=", parent_id_);
-  }
-
+  string result = strings::StrCat(params_.prefix, "#", traceme_metadata_);
   TraceMeMetadata metadata = GetTraceMeMetadata();
   for (const auto& pair : metadata) {
     strings::StrAppend(&result, ",", pair.first, "=", pair.second);
diff --git a/tensorflow/core/framework/dataset.h b/tensorflow/core/framework/dataset.h
index a97142e8cf2..1ac3c77d8d8 100644
--- a/tensorflow/core/framework/dataset.h
+++ b/tensorflow/core/framework/dataset.h
@@ -1094,6 +1094,7 @@ class DatasetBaseIterator : public IteratorBase {
     return model && model->collect_resource_usage() && node_;
   }
 
+  string traceme_metadata_;
   BaseParams params_;
 };