From bfafc1acef59ff5a7ba2bf2675350812e552d5ad Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 18 Mar 2020 14:38:13 -0700
Subject: [PATCH] don't trace arguments (include tensor shapes and op
 attributes used for cost analysis) for dataset ops.

PiperOrigin-RevId: 301671357
Change-Id: I4cf44855603ea26007a9652d6c01866db4d83c5b
---
 tensorflow/core/framework/dataset.cc | 4 ++++
 tensorflow/core/framework/dataset.h  | 2 ++
 2 files changed, 6 insertions(+)

diff --git a/tensorflow/core/framework/dataset.cc b/tensorflow/core/framework/dataset.cc
index 97c4d212223..cccbdd5d8e4 100644
--- a/tensorflow/core/framework/dataset.cc
+++ b/tensorflow/core/framework/dataset.cc
@@ -508,6 +508,10 @@ void DatasetOpKernel::Compute(OpKernelContext* ctx) {
   }
 }
 
+string DatasetOpKernel::TraceString(OpKernelContext* ctx, bool verbose) {
+  return strings::StrCat(name_view(), ":", type_string_view());
+}
+
 // static
 bool DatasetOpKernel::IsDatasetOp(const OpDef* op_def) {
   if (DatasetOpRegistry::IsRegistered(op_def->name())) {
diff --git a/tensorflow/core/framework/dataset.h b/tensorflow/core/framework/dataset.h
index 92f7a52b632..25cc8fd759e 100644
--- a/tensorflow/core/framework/dataset.h
+++ b/tensorflow/core/framework/dataset.h
@@ -1073,6 +1073,8 @@ class DatasetOpKernel : public OpKernel {
   // the `DatasetOpKernel` class.
   static bool IsDatasetOp(const OpDef* op_def);
 
+  string TraceString(OpKernelContext* ctx, bool verbose) override;
+
  protected:
   // Subclasses should implement this method. It will be called during Compute
   // execution.