message: Replace existing shape inference function if it exists.
diff: @@ -394,8 +394,8 @@ Status ConstructHostGraph(
  for (const string& host_func : outside_compilation_host_graphs) {
  VLOG(4) << "Expanding host graph " << host_func;
  FunctionBody* host_fbody = nullptr;
  - TF_RETURN_IF_ERROR(
  - FunctionDefToBodyHelper(*fld->Find(host_func), AttrSlice(), fld,
  + TF_RETURN_IF_ERROR(FunctionDefToBodyHelper(
  + *fld->Find(host_func), AttrSlice(), fld,
  [&](const string& op, const OpDef** sig) {
  return fld->LookUpOpDef(op, sig);
  },
  @@ -411,7 +411,8 @@ Status ConstructHostGraph(
  node_map[host_fbody->graph->source_node()] = (*host_graph)->source_node();
  node_map[host_fbody->graph->sink_node()] = (*host_graph)->sink_node();
  Status s;
  - ReverseDFS(*host_fbody->graph, /*enter=*/nullptr,
  + ReverseDFS(
  + *host_fbody->graph, /*enter=*/nullptr,
  [&](const Node* n) {
  if (!s.ok()) {
  return;
  @@ -838,10 +839,15 @@ Status ExtractOutsideCompilationForFunction(
  FunctionDef shape_inference_fdef = *xla_fdef;
  shape_inference_fdef.mutable_signature()->set_name(
  shape_inference_graph);
  + if (fld->Find(shape_inference_graph)) {
  + TF_RETURN_IF_ERROR(fld->ReplaceFunction(shape_inference_graph,
  + shape_inference_fdef));
  + } else {
  TF_RETURN_IF_ERROR(fld->AddFunctionDef(shape_inference_fdef));
  }
  }
  }
  + }
  for (Node* n : outside_compilation_nodes) {
  TF_RETURN_IF_ERROR(ReplaceOrRemoveOutsideCompilationCallNode(
  graph_out.get(), n, host_compute_core));
  