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

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

 private:
  LogicalResult matchAndRewrite(TF::PartitionedCallOp call_op,
                                PatternRewriter& rewriter) const override {
    StringRef function_name =
        mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
    if (!function_name.starts_with("composite_") ||
        !call_op->hasAttr(kQuantTraitAttrName)) {
      return failure();
    }

      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

  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 =

    if (method.ok() && method->has_static_range_ptq()) return true;
  }

  TF::PartitionedCallOp call_op = dyn_cast_or_null<TF::PartitionedCallOp>(op);
  return call_op && call_op->hasAttrOfType<StringAttr>(kQuantTraitAttrName) &&
         call_op->getAttrOfType<StringAttr>(kQuantTraitAttrName).getValue() ==
             llvm::StringRef(

  ASSERT_TRUE(module_op);

  func::FuncOp main_fn = FindMainFuncOp(*module_op);
  ASSERT_THAT(main_fn, NotNull());

  auto call_op = *main_fn.getOps<TF::XlaCallModuleOp>().begin();
  EXPECT_TRUE(HasWeightOnlyPtqMethod(call_op));
}

TEST_F(LiftAsFunctionCallTest, HasWeightOnlyPtqMethodDifferentMethod) {
  const absl::string_view kModuleDotNoQuantization = R"mlir(
    module {

  return method.has_weight_only_ptq();
}

bool IsWeightOnlyQuantizableOp(const Operation& op) {
  if (auto call_op = dyn_cast<TF::XlaCallModuleOp>(op)) {
    StringRef entry_function_name = GetEntryFunctionName(call_op);
    absl::StatusOr<Method> quantization_method = GetQuantizationMethod(call_op);
    return ContainsConvOrDot(entry_function_name) && quantization_method.ok() &&
           quantization_method->has_weight_only_ptq();

  bool IsCompositeFunction(TF::PartitionedCallOp call_op) const {
    if (!call_op->hasAttr(kQuantTraitAttrName)) {
      return false;
    }

    const auto f_attr = call_op.getFAttr().dyn_cast<FlatSymbolRefAttr>();
    if (!f_attr || !f_attr.getValue().starts_with("composite_")) {
      return false;
    }

    bool has_quantized_types = false;
    for (Value input : call_op.getArgs()) {

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