: public OpRewritePattern<func::CallOp> {
 public:
  explicit MergeFusionWithUniformDequantizePattern(MLIRContext* context)
      : OpRewritePattern<func::CallOp>(context) {}
  LogicalResult matchAndRewrite(func::CallOp call_op,
                                PatternRewriter& rewriter) const override {
    if (call_op.getNumResults() != 1) return failure();
    auto users = call_op->getUsers();
    for (auto user : users) {

};

void EncapsulatePartitionedCall(Operation *call_op,
                                mlir::StringAttr callee_name) {
  OpBuilder builder(call_op);
  auto cluster = builder.create<mlir::tf_device::ClusterOp>(
      call_op->getLoc(), call_op->getResultTypes());
  cluster.getBody().push_back(new Block);
  call_op->replaceAllUsesWith(cluster.getResults());
  call_op->moveBefore(&cluster.GetBody(), cluster.GetBody().end());

std::optional<QuantizationResult> GetQuantizationResult(func::CallOp call_op) {
  const StringRef callee_name = call_op.getCalleeAttr().getValue();
  if (!callee_name.starts_with(kQuantizedFuncPrefix)) {
    return std::nullopt;  // `call_op` is not a quantized function call.
  }

  absl::StatusOr<Method> method = GetQuantizationMethod(call_op);
  if (!method.ok()) {
    call_op->emitError() << "Failed to get quantization method: "

  // for the call_op.
  Operation* last_output = subgraph.partition_ops_.back();

  builder.setInsertionPoint(last_output);
  auto call_op = builder.create<func::CallOp>(last_output->getLoc(), func,
                                              subgraph.FuncArguments());
  ops_added.call_op = call_op;
  // FuncOutputs refer to the original `Values` in input module which are now

void RestoreFunctionNameFromXlaCallModuleOp(TF::XlaCallModuleOp& call_op,
                                            SymbolTable& symbol_table) {
  if (!call_op->hasAttr(kOriginalStablehloEntryFunctionAttrName)) {
    return;
  }

  const auto original_function_name = call_op->getAttrOfType<StringAttr>(
      kOriginalStablehloEntryFunctionAttrName);
  const auto current_function_name = call_op->getAttrOfType<FlatSymbolRefAttr>(

    PatternRewriter& rewriter, CallOp call_op,
    const SmallVectorImpl<Type>& output_types,
    const SmallVectorImpl<Value>& inputs, const NamedAttrList& attr_list,
    const llvm::StringMap<Attribute>& derived_attrs) const {
  // Create the new op
  Location loc = call_op.getLoc();
  rewriter.setInsertionPointAfter(call_op);
  std::string tf_op_name = GetTFOpName(call_op.getCallee());

};

void UnwrapXlaCallModuleOp(TF::XlaCallModuleOp call_op,
                           SymbolTable& symbol_table) {
  // Do not inline lifted quantized functions used for fusing patterns.
  // TODO - b/310539922: Remove reference to TF/TFL utils.
  if (call_op->hasAttr(kQuantTraitAttrName)) {
    return;
  }

  auto function_name = call_op
                           ->getAttrOfType<FlatSymbolRefAttr>(

bool IsQuantizedCallforDynamicRange(TF::PartitionedCallOp call_op) {
  bool has_quantized_types_for_weights = false;
  std::unique_ptr<OpQuantSpec> spec = GetTFOpQuantSpec(call_op);

  for (int32_t cur_idx = 0; cur_idx < call_op.getArgs().size(); cur_idx++) {
    // Check if the only the weight index has QuantizeCastOp.
    auto cur_op = dyn_cast_or_null<quantfork::QuantizeCastOp>(
        call_op.getArgs()[cur_idx].getDefiningOp());

    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(
    TF::XlaCallModuleOp call_op, const FlatSymbolRefAttr &f_attr) {

  using OpRewritePattern<TF::TPUPartitionedCallOp>::OpRewritePattern;

 private:
  LogicalResult matchAndRewrite(TF::TPUPartitionedCallOp call_op,
                                PatternRewriter& rewriter) const override {
    auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
    auto module_op = call_op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module_op);

    auto f_name = f_attr.getValue();