func(node_def);
  }

  for (tensorflow::FunctionDef& function_def :
       *graph_def.mutable_library()->mutable_function()) {
    for (tensorflow::NodeDef& node_def : *function_def.mutable_node_def()) {
      func(node_def);
    }
  }
}

}  // namespace stablehlo::quantization

  if (library) {
    // Upgrade nodes in the functions.
    for (FunctionDef& fdef : *library->mutable_function()) {
      auto func_name = fdef.signature().name();
      for (auto& node_def : *fdef.mutable_node_def()) {
        const OpDef* op_def = nullptr;
        // With lazy loading, some functions might not be executed, thus we skip
        // the node if the op is not registered.

  // Remove "_xla_inferred_shapes" attr. They are added by
  // `PerformStaticShapeInferenceBeforeEncapsulation`.
  for (FunctionDef& fdef : *library->mutable_function()) {
    for (NodeDef& node_def : *fdef.mutable_node_def()) {
      node_def.mutable_attr()->erase("_xla_inferred_shapes");
    }
  }

  return s;
}

Status Encapsulate(GraphDef* graphdef, FunctionDefLibrary* library) {