| 22 | namespace grappler { |
| 23 | |
| 24 | bool IsTPUGraphDef(const GraphDef& def) { |
| 25 | for (const auto& node : def.node()) { |
| 26 | if (node.op() == "TPUCompile" || node.op() == "TPUPartitionedCall") { |
| 27 | return true; |
| 28 | } |
| 29 | } |
| 30 | if (!def.has_library()) return false; |
| 31 | for (const auto& function_def : def.library().function()) { |
| 32 | for (const auto& node : function_def.node_def()) { |
| 33 | if (node.op() == "TPUCompile" || node.op() == "TPUPartitionedCall") { |
| 34 | return true; |
| 35 | } |
| 36 | } |
| 37 | } |
| 38 | return false; |
| 39 | } |
| 40 | |
| 41 | } // end namespace grappler |
| 42 | } // end namespace tensorflow |