| 21 | namespace tensorflow { |
| 22 | |
| 23 | Status ShapeAnnotationsMatch( |
| 24 | const Graph& graph, const GraphShapeInfo& shape_info, |
| 25 | std::map<string, std::vector<PartialTensorShape>> expected_shapes) { |
| 26 | for (Node* node : graph.op_nodes()) { |
| 27 | auto sit = shape_info.find(node->name()); |
| 28 | TF_RET_CHECK(sit != shape_info.end()) |
| 29 | << "Missing shape information for node " << node->name(); |
| 30 | std::vector<PartialTensorShape> shapes; |
| 31 | for (const auto& output : sit->second) shapes.push_back(output.shape); |
| 32 | |
| 33 | auto it = expected_shapes.find(node->name()); |
| 34 | if (it != expected_shapes.end()) { |
| 35 | if (!PartialTensorShapeUtils::AreIdentical(shapes, it->second)) { |
| 36 | return errors::InvalidArgument( |
| 37 | "Shape mismatch for ", node->name(), ". Expected: ", |
| 38 | PartialTensorShapeUtils::PartialShapeListString(it->second), |
| 39 | ", actual: ", |
| 40 | PartialTensorShapeUtils::PartialShapeListString(shapes)); |
| 41 | } |
| 42 | expected_shapes.erase(it); |
| 43 | } |
| 44 | } |
| 45 | if (!expected_shapes.empty()) { |
| 46 | std::vector<string> missing; |
| 47 | missing.reserve(expected_shapes.size()); |
| 48 | for (const auto& entry : expected_shapes) { |
| 49 | missing.push_back(entry.first); |
| 50 | } |
| 51 | return errors::InvalidArgument("Missing shapes for nodes: ", |
| 52 | absl::StrJoin(missing, ",")); |
| 53 | } |
| 54 | return Status::OK(); |
| 55 | } |
| 56 | |
| 57 | } // namespace tensorflow |