return GetInverseScalesConstantOp().getValue();
  }

  func::CallOp GetCallOp() { return call_op_; }

  FlatSymbolRefAttr GetFunction() { return call_op_.getCalleeAttr(); }

 private:
  explicit UniformQuantizeFunctionCallPattern(func::CallOp call_op)
      : call_op_(call_op) {}

  func::CallOp call_op_;
};

// Matches the following pattern that represents uniform dequantization.
//

};

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

  // 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

    : 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) {

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: "

      aggregated_cost_with_decisions;
};

// If the output is produced by a callop, will return the callop, otherwise,
// will return nullptr.
inline func::CallOp GetProducerCallOpOrNull(Value output) {
  Operation* output_op = output.getDefiningOp();
  if (output_op != nullptr && llvm::isa<func::CallOp>(output_op)) {
    return llvm::cast<func::CallOp>(output_op);
  }
  return nullptr;
}

class PickSubgraphsPass

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

// Checks if an op calls a composite function and all the inputs and outputs are
// quantized.
bool IsQuantizedCompositeFunction(func::CallOp call_op) {
  if (!call_op.getCallee().starts_with("quantized_")) {
    return false;
  }

  bool has_quantized_types = false;
  for (Value operand : call_op.getOperands()) {
    if (const TensorType type = mlir::dyn_cast<TensorType>(operand.getType())) {

  bool changed = false;
  auto walk_result = func.walk([&](CallOp call_op) {
    auto callee = table.lookup<TFRFuncOp>(call_op.getCallee());
    if (!callee || callee.isExternal()) return WalkResult::advance();

    // Record the boundary of the inlined operations. The inlined operation will
    // be inserted between these two operations.
    Operation* inlined_point = call_op.getOperation();
    Operation* after_inlined_point =

  Block* then_block = builder.createBlock(merge_block);
  Operation* call_op = CallFn(loc, get_operand, op.then_function(), &builder);

  auto get_then_result = [&](int i) { return call_op->getResult(i); };
  JumpToBlock(loc, get_then_result, merge_block, &builder);

  // Set up the 'else' block.
  Block* else_block = builder.createBlock(merge_block);
  call_op = CallFn(loc, get_operand, op.else_function(), &builder);