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);