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 {

  for (auto& function_use : function_uses.value()) {
    auto call_op = function_use.getUser();
    auto caller_func = call_op->getParentOfType<func::FuncOp>();
    if (!caller_func) return failure();

    builder.setInsertionPoint(call_op);
    for (auto [lifted_op, arg_idx] : lifted_op_and_arg_idx) {
      auto new_op = builder.clone(*lifted_op, mapping);
      call_op->insertOperands(arg_idx, new_op->getResult(0));
    }

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

    }
  }
}

}  // namespace

std::unique_ptr<OpQuantSpec> GetStableHloOpQuantSpec(Operation* op) {
  auto spec = std::make_unique<OpQuantSpec>();
  if (auto call_op = dyn_cast_or_null<TF::XlaCallModuleOp>(op)) {
    auto entry_function =
        call_op->getAttrOfType<FlatSymbolRefAttr>("_entry_function");
    StringRef function_name = entry_function.getValue();
    if (!function_name.starts_with("composite_")) {
      return spec;

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

// quantization.
inline FlatSymbolRefAttr GetFuncAttr(TF::PartitionedCallOp call_op) {
  return mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
}

inline FlatSymbolRefAttr GetFuncAttr(TF::XlaCallModuleOp call_op) {
  return call_op->getAttrOfType<FlatSymbolRefAttr>(
      TF::kStablehloEntryFunctionAttrName);
}

      }
      if (failed(HandleCaseOrIfOp(case_op, branch_functions))) return failure();
    } else if (auto call_op = llvm::dyn_cast<TF::PartitionedCallOp>(&op)) {
      auto callee = call_op.func();
      if (!callee) {
        return call_op.emitOpError(
            "resource lifting does not support call with nested references.");
      }
      if (failed(HandlePartitionedCallOp(call_op, callee, module,

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

  Operation* call_op = FindOperationOfType<TF::XlaCallModuleOp>(main_fn);
  EXPECT_FALSE(IsHybridQuantizedOp(call_op));
}

constexpr absl::string_view kModuleDotGeneralFullyConnected = R"mlir(
  module {

    if (!main_func) continue;

    SymbolTable symbol_table(module_op);
    for (auto call_op : main_func.getOps<TF::PartitionedCallOp>()) {
      func_ops.push_back(dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(
          mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue())));
    }
    for (auto call_op : main_func.getOps<TF::StatefulPartitionedCallOp>()) {
      func_ops.push_back(