Function FindDynamicStitchNode

serving/processor/framework/graph_optimizer.cc:968–1003

}

Node* FindDynamicStitchNode(Node* gather_node) {
  // Supported pattern:
  //
  //        KvResourceGather
  //              |
  //           Identity
  //              |
  //             ...
  //              |
  //           Identity
  //              |
  //     ParallelDynamicStitch
  //
  if (gather_node->num_outputs() != 1) {
    LOG(FATAL) << "Gather node has more than one edge, "
               << gather_node->DebugString();
  }

  for (const Edge* edge : gather_node->out_edges()) {
    Node* dest = edge->dst();
    while (dest->op_def().name() == "Identity") {
      Node* tmp = nullptr;
      for (const Edge* e : dest->out_edges()) {
        tmp = e->dst();
        break;
      }
      dest = tmp;
    }

    if (IsDynamicStitchOp(dest)) {
      return dest;
    }
  }

  return nullptr;
}

Node* CheckDynamicStitchNode(std::vector<Node*> gather_nodes) {
  Node* ds_node = nullptr;