std::optional<QuantizationUnitLoc::QuantizationUnit> found_unit =
          FindQuantizationUnitFromLoc(child_loc);
      if (found_unit.has_value()) return found_unit;
    }
  } else if (isa<CallSiteLoc>(loc)) {
    // If the graph is inlined, CallSiteLoc can be created.
    return FindQuantizationUnitFromLoc(
        mlir::cast<CallSiteLoc>(loc).getCallee());
  }

  return std::nullopt;
}

  static bool classof(Attribute attr);
};

// Finds the QuantizationUnit from location info.
std::optional<QuantizationUnitLoc::QuantizationUnit>
FindQuantizationUnitFromLoc(Location loc);

}  // namespace quant
}  // namespace mlir

 private:
  LogicalResult matchAndRewrite(Operation* op,
                                PatternRewriter& rewriter) const override {
    if (!IsOpWithQuantizableTrait(op) ||
        FindQuantizationUnitFromLoc(op->getLoc()).has_value()) {
      return failure();
    }

    std::optional<QuantizationUnit> quantization_unit =
        FindQuantizationUnit(op);

std::pair<std::string, std::string> GetFuncNameAndNodeName(
    TF::PartitionedCallOp call_op, const FlatSymbolRefAttr &f_attr) {
  std::optional<QuantizationUnitLoc::QuantizationUnit> quant_unit =
      FindQuantizationUnitFromLoc(call_op->getLoc());
  return std::make_pair(quant_unit->func_name(), quant_unit->node_name());
}

std::pair<std::string, std::string> GetFuncNameAndNodeName(

  std::optional<QuantizationMethod> getUnitWiseQuantizationMethod(
      TF::PartitionedCallOp call_op) const {
    // If unit-wise quantization config is found, overwrite the default config.
    auto quantization_unit = FindQuantizationUnitFromLoc(call_op.getLoc());
    if (!quantization_unit.has_value()) return std::nullopt;

    for (const auto& unit_config :
         quant_options_.unit_wise_quantization_specs()) {

  // Set the location of call op to QuantizationUnitLoc if found.
  Location call_op_loc = location;
  for (Operation* op : cloning_ops) {
    std::optional<QuantizationUnitLoc::QuantizationUnit> unit =
        FindQuantizationUnitFromLoc(op->getLoc());
    if (unit.has_value()) {
      call_op_loc = QuantizationUnitLoc(builder.getContext(), unit.value());
    }
  }

  if (failed(SetAttributeMap(*context, attributes, cloning_ops))) {