| 563 | } |
| 564 | |
| 565 | Status Encapsulator::Subgraph::BuildFunctionDef( |
| 566 | const string& name_in, const RewriteSubgraphFn& rewrite_subgraph_fn, |
| 567 | bool reuse_existing_functions, FunctionLibraryDefinition* library) { |
| 568 | // name_in is copied here because name may be modified below if |
| 569 | // rewrite_subgraph_fn is true. |
| 570 | string name = name_in; |
| 571 | call_node_def_.set_op(name); |
| 572 | call_node_def_.set_name(name); |
| 573 | call_node_def_.set_device(device_); |
| 574 | |
| 575 | if (rewrite_subgraph_fn) { |
| 576 | std::vector<OutputTensor> arg_source_tensors(args_by_src_.size()); |
| 577 | for (const auto& arg : args_by_src_) { |
| 578 | arg_source_tensors.at(arg.second) = arg.first; |
| 579 | } |
| 580 | // Initialize the input and output permutations to the identity. |
| 581 | std::vector<int> input_permutation(args_by_src_.size()); |
| 582 | std::iota(input_permutation.begin(), input_permutation.end(), 0); |
| 583 | std::vector<int> output_permutation(results_.size()); |
| 584 | std::iota(output_permutation.begin(), output_permutation.end(), 0); |
| 585 | |
| 586 | TF_RETURN_IF_ERROR( |
| 587 | rewrite_subgraph_fn(arg_source_tensors, &graph_, &input_permutation, |
| 588 | &output_permutation, &call_node_def_)); |
| 589 | |
| 590 | // Apply the input/output permutations to the 'args_by_...' and 'results_' |
| 591 | // mappings, so when we build edges in BuildOutputGraph() we |
| 592 | // connect them to the right input/output positions. |
| 593 | if (input_permutation.size() != args_by_src_.size()) { |
| 594 | return errors::InvalidArgument("Input permutation has incorrect size."); |
| 595 | } |
| 596 | if (output_permutation.size() != results_.size()) { |
| 597 | return errors::InvalidArgument("Output permutation has incorrect size."); |
| 598 | } |
| 599 | for (auto& arg : args_by_src_) { |
| 600 | arg.second = input_permutation[arg.second]; |
| 601 | } |
| 602 | for (auto& arg : args_by_dst_) { |
| 603 | arg.second = input_permutation[arg.second]; |
| 604 | } |
| 605 | for (auto& result : results_) { |
| 606 | result.second = output_permutation[result.second]; |
| 607 | } |
| 608 | |
| 609 | name = call_node_def_.op(); |
| 610 | } |
| 611 | |
| 612 | function_def_name_ = name; |
| 613 | |
| 614 | FunctionDef fdef; |
| 615 | auto lookup = [this](const Node* node) -> absl::optional<string> { |
| 616 | if (control_output_nodes_.contains(node->name())) { |
| 617 | return absl::make_optional(node->name()); |
| 618 | } |
| 619 | return absl::nullopt; |
| 620 | }; |
| 621 | // Verify that the graph has well-formed control flow structure. |
| 622 | std::vector<ControlFlowInfo> dummy; |
no test coverage detected