| 593 | } |
| 594 | |
| 595 | Status SavedModelOptimizer::RunNativeTFGraphPass() { |
| 596 | if (option_.shard_embedding) { |
| 597 | // Find all variables ops which should be rewrite, |
| 598 | // include partitioned or non-partitoned variables |
| 599 | std::unordered_map<std::string, std::vector<Node*>> var_parts; |
| 600 | TF_RETURN_IF_ERROR(FindVariableParts(var_parts)); |
| 601 | |
| 602 | std::unordered_map<std::string, std::vector<Node*>> import_nodes; |
| 603 | TF_RETURN_IF_ERROR(FindKvResourceImportNode(var_parts, import_nodes)); |
| 604 | |
| 605 | // Rewrite subgraph |
| 606 | TF_RETURN_IF_ERROR(RewriteEmbeddingLookupGraph(var_parts, import_nodes)); |
| 607 | } |
| 608 | |
| 609 | if (option_.st) { |
| 610 | TF_RETURN_IF_ERROR(RewriteEmbeddingVariableAttr(option_.st, |
| 611 | option_.path, option_.size)); |
| 612 | } |
| 613 | |
| 614 | // Add other passes here |
| 615 | |
| 616 | // replace the graph def in saved_model_bundle |
| 617 | graph_.ToGraphDef(meta_graph_def_->mutable_graph_def()); |
| 618 | |
| 619 | return Status::OK(); |
| 620 | } |
| 621 | |
| 622 | Status SavedModelOptimizer::RunODLGraphPass() { |
| 623 | // Generate ids for every feature |
nothing calls this directly
no test coverage detected