TF_ASSERT_OK(CreateMlirModule("basic_cluster.mlir"));

  TF_EXPECT_OK(RunLowerClusterToRuntimeOpsPassPipeline(
      *mlir_module_, DeviceType(DEVICE_TPU_XLA_JIT)));

  FuncOp main = mlir_module_->lookupSymbol<FuncOp>("main");
  ASSERT_TRUE(main);

  bool has_cluster_op = false;
  main.walk([&](mlir::tf_device::ClusterOp) {
    has_cluster_op = true;
    return mlir::WalkResult::interrupt();
  });

)mlir";

TEST_F(IsOpQuantizableStableHloTest, ConstantOpQuantizable) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleConstantAdd);
  ASSERT_TRUE(module_op);

  auto test_func = module_op->lookupSymbol<func::FuncOp>("constant_add");
  ASSERT_THAT(test_func, NotNull());

  auto constant_op =
      FindOperationOfType<mlir::stablehlo::ConstantOp>(test_func);
  EXPECT_TRUE(IsOpQuantizableStableHlo(constant_op));
}

    // restriction, it would require tensor cast ops.
    func::FuncOp parent_op = op->getParentOfType<func::FuncOp>();
    if (!llvm::hasSingleElement(parent_op)) return failure();

    // Find the then and else branch functions.
    func::FuncOp then_func = op.then_function();
    func::FuncOp else_func = op.else_function();

    // If the If has no uses and its functions are side-effect free, then

// Create an instance of the TargetAnnotationPass.
// TODO(b/177376459): Remove in favor of the one below.
std::unique_ptr<OperationPass<func::FuncOp>> CreateTargetAnnotationPass(
    llvm::ArrayRef<std::string> device_specs);

// Create and instance of TargetAnnotationPass.
std::unique_ptr<OperationPass<func::FuncOp>> CreateTargetAnnotationPass(
    const TacModule* module);

// Create an instance of the RaiseTargetSubgraphsPass. If `skip_raise_cpu_ops`,

  void ConvertTFImplements(func::FuncOp func, StringAttr attr);
  void ConvertTFImplementsWithAttributes(func::FuncOp func, FuncAttr attr);
  void ConvertTFAPIImplements(func::FuncOp func, StringAttr attr,
                              ModuleOp module);
  void runOnOperation() override;
};

LogicalResult CheckFusableLayerNormalizedLstmCellSimple(
    func::FuncOp lstm_func) {
  for (int i = 0; i < 5; ++i) {

//
// The current implementation follows direct calls to `mlir::func::FuncOp` only
// and returns a `mlir::WalkResult::interrupt()` when it encounters a call whose
// callee cannot be resolved to `mlir::func::FuncOp`.
mlir::WalkResult WalkReachableFunctions(
    mlir::func::FuncOp func,
    llvm::function_ref<mlir::WalkResult(mlir::func::FuncOp)> callback,
    mlir::SymbolTableCollection* symbol_table = nullptr);

  // Converts a given FuncOp to a FunctionDef and adds it to the function
  // definition library
  static Status ConvertLibFunction(
      const GraphExportConfig& configs, const Dialect* tf_dialect,
      const SymbolTable& symbol_table, FuncOp function,
      FunctionLibraryDefinition* flib_def,
      llvm::SmallDenseSet<FuncOp>& visited_functions);

  // Converts a given FuncOp to a FunctionDef and adds it to the function
  // definition library
  static Status ConvertLibFunction(
      const GraphExportConfig& configs, const Dialect* tf_dialect,
      const SymbolTable& symbol_table, FuncOp function,
      FunctionLibraryDefinition* flib_def,
      llvm::SmallDenseSet<FuncOp>& visited_functions);

  FunctionType func_type =
      FunctionType::get(rewriter.getContext(), {create_unknown_input_shape},
                        {create_unknown_output_shape});

  func::FuncOp quantization_func =
      rewriter.create<func::FuncOp>(input_op->getLoc(), func_name, func_type);

  OpBuilder::InsertionGuard guard = OpBuilder::InsertionGuard(rewriter);
  ArrayRef<Type> inputs = quantization_func.getFunctionType().getInputs();

  std::string function_name = absl::StrCat("func_", subgraph.subgraph_id_);

  func::FuncOp new_func = func::FuncOp::create(builder.getUnknownLoc(),
                                               function_name, function_type);
  new_func.setVisibility(func::FuncOp::Visibility::Private);
  new_func.addEntryBlock();

  // To form the body of the new function we need to clone each