Disable PinToHostOptimizer for any TPU graphs.

PiperOrigin-RevId: 214338297

1 files changed, 15 insertions, 0 deletions

diff --git a/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc b/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
index c8f9311b2e..98c27300a9 100644
--- a/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/pin_to_host_optimizer.cc
@@ -157,6 +157,16 @@ string TryFindHostDevice(const gtl::FlatSet<string>& devices,
   return device;
 }
 
+bool IsTPUGraphDef(const GraphDef& def) {
+  for (auto node : def.node()) {
+    if (node.op() == "TPUCompile" || node.op() == "TPUExecute" ||
+        node.op() == "TPUPartitionedCall") {
+      return true;
+    }
+  }
+  return false;
+}
+
 // All the nodes that should be blacklisted and not swapped.
 bool IsBlacklisted(const NodeDef& node) { return IsCollective(node); }
 }  // end namespace internal
@@ -165,6 +175,11 @@ Status PinToHostOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
                                     GraphDef* optimized_graph) {
   *optimized_graph = item.graph;
 
+  // Skip all TPU graphs.
+  if (internal::IsTPUGraphDef(*optimized_graph)) {
+    return Status::OK();
+  }
+
   GraphProperties properties(item);
   bool has_properties = false;
   GraphView graph(optimized_graph);