| 810 | } |
| 811 | |
| 812 | Status Conditional::AddInputEdges( |
| 813 | Graph* graph, |
| 814 | const std::unordered_map<Node*, OutputTensor>& merge_to_replacement) { |
| 815 | VLOG(2) << "AddInputEdges for " << if_node_->name(); |
| 816 | int index = 0; |
| 817 | // Add predicate input. |
| 818 | if (predicate_.node->IsMerge()) { |
| 819 | // If the predicate is a Merge node, we should not use Merge output as |
| 820 | // predicate. Instead, we should use the corresponding If output in |
| 821 | // 'merge_to_replacement'. Otherwise, this Conditional's If node is still |
| 822 | // connected to the predicate Merge node; and when we call |
| 823 | // DeleteReachableAndDeadNodes(), the predicate Merge node and this |
| 824 | // Conditional's If node will be removed. |
| 825 | auto iter = merge_to_replacement.find(predicate_.node); |
| 826 | if (iter == merge_to_replacement.end()) { |
| 827 | return errors::Internal("Cannot find replacement for Merge node ", |
| 828 | predicate_.node->name()); |
| 829 | } |
| 830 | graph->AddEdge(iter->second.node, iter->second.index, if_node_, index++); |
| 831 | } else { |
| 832 | graph->AddEdge(const_cast<Node*>(predicate_.node), predicate_.index, |
| 833 | if_node_, index++); |
| 834 | } |
| 835 | // Add function body inputs. |
| 836 | for (auto& arg : cond_arg_nodes_) { |
| 837 | if (arg.src_output == Graph::kControlSlot) { |
| 838 | graph->AddControlEdge(arg.src, if_node_); |
| 839 | } else { |
| 840 | graph->AddEdge(arg.src, arg.src_output, if_node_, index++); |
| 841 | } |
| 842 | } |
| 843 | for (Node* n : external_control_inputs_) { |
| 844 | graph->AddControlEdge(n, if_node_); |
| 845 | } |
| 846 | return Status::OK(); |
| 847 | } |
| 848 | |
| 849 | Status Conditional::AddOutputEdges( |
| 850 | Graph* graph, |