| 464 | Graph* Encapsulator::Subgraph::GetGraph() const { return graph_.get(); } |
| 465 | |
| 466 | Status Encapsulator::Subgraph::RecordArg( |
| 467 | const Edge* edge, const std::unordered_map<const Node*, Node*>& node_images, |
| 468 | std::vector<std::pair<const Node*, Node*>>* src_arg_pairs) { |
| 469 | Node* src_node = edge->src(); |
| 470 | int src_slot = edge->src_output(); |
| 471 | std::unordered_map<OutputTensor, int, OutputTensor::Hash>::iterator iter; |
| 472 | bool inserted; |
| 473 | std::tie(iter, inserted) = args_by_src_.emplace( |
| 474 | OutputTensor(src_node, src_slot), args_by_src_.size()); |
| 475 | int arg_index = iter->second; |
| 476 | if (inserted) { |
| 477 | NodeDef arg_def; |
| 478 | NodeDefBuilder builder( |
| 479 | absl::StrCat(src_node->name(), "_", src_slot, "_arg"), kArgOp, |
| 480 | NodeDebugInfo(src_node->def())); |
| 481 | DataType dtype = edge->dst()->input_type(edge->dst_input()); |
| 482 | builder.Attr("T", dtype); |
| 483 | builder.Attr("index", arg_index); |
| 484 | Status s = builder.Finalize(&arg_def); |
| 485 | if (!s.ok()) return s; |
| 486 | |
| 487 | Node* arg = graph_->AddNode(arg_def, &s); |
| 488 | if (!s.ok()) return s; |
| 489 | |
| 490 | src_arg_pairs->push_back({src_node, arg}); |
| 491 | args_.push_back(arg); |
| 492 | } |
| 493 | Node* dst_node = edge->dst(); |
| 494 | Node* dst_image = node_images.at(dst_node); |
| 495 | int dst_slot = edge->dst_input(); |
| 496 | args_by_dst_[InputTensor(dst_node, dst_slot)] = arg_index; |
| 497 | graph_->AddEdge(args_[arg_index], 0, dst_image, dst_slot); |
| 498 | return Status::OK(); |
| 499 | } |
| 500 | |
| 501 | Status Encapsulator::Subgraph::RecordControlResult( |
| 502 | const Edge* edge, |
no test coverage detected