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

func::FuncOp FindMainFuncOp(ModuleOp module_op) {
  if (const auto main_func_op = module_op.lookupSymbol<func::FuncOp>(
          kImportModelDefaultGraphFuncName);
      main_func_op != nullptr && IsPublicFuncOp(main_func_op)) {
    return main_func_op;
  }

  if (const auto serving_default_func_op =
          module_op.lookupSymbol<func::FuncOp>(kDefaultServingSignatureDefKey);
      serving_default_func_op != nullptr &&

    }
  )mlir";

  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kFullyQuantizedAdd);
  ASSERT_TRUE(module_op);

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

  auto add_op_itr = func_op.getBody().op_begin<mlir::stablehlo::AddOp>();
  ASSERT_THAT(add_op_itr,

        func.return %mul : tensor<1024x24x24x3xf32>
      }
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op_ref = ParseModuleOpString(kModuleCode);
  const auto test_func =
      module_op_ref->lookupSymbol<func::FuncOp>("test_fold_constant");
  ASSERT_THAT(test_func, NotNull());

  Operation* mul_op = FindOperationOfType<TF::MulOp>(test_func);
  SmallVector<Value> results = ConstantFoldOpIfPossible(mul_op);

  TF_EXPECT_OK(RunFunctionTf2xlaClusteringBridge(
      *mlir_module_, /*is_supported_by_replicated_brige*/ true,
      /*is_in_fallback_enabled_mode=*/false));

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

  EXPECT_EQ(compilation_status.Delta(mlir::TF::kMlirPh1BridgeCounterReplicated,
                                     mlir::TF::kMlirPh1BridgeCounterV2,

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

absl::Status SetupArguments(mlir::ModuleOp module,
                            std::vector<TensorShape>& arg_shapes,
                            tpu::TPUCompileMetadataProto& metadata_proto) {
  auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
  if (!main_fn) {
    return absl::InternalError("Could not find main function in MLIR Module.");
  }

  mlir::FunctionType func_type = main_fn.getFunctionType();

  if (!status.ok()) {
    return op.emitError()
           << "attribute 'host_mlir_module' can not be deserialized. "
           << status.message();
  }

  func::FuncOp func = module_for_func->lookupSymbol<func::FuncOp>("host_func");
  if (!func)
    return op.emitError()
           << "serialized module in attribute 'host_mlir_module' does not "
              "contain 'host_func' function.";

    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
    std::vector<xla::Shape>* xla_input_shapes) {
  xla_input_shapes->clear();

  mlir::func::FuncOp main_func =
      module.lookupSymbol<mlir::func::FuncOp>("main");
  TF_RET_CHECK(main_func != nullptr) << "No main function found";
  mlir::FunctionType func_type = main_func.getFunctionType();

  int num_args = func_type.getNumInputs();

        mirrored_replicate_args.insert(
            mlir::cast<IntegerAttr>(mirrored_index).getInt());
      }
    }
    auto func =
        llvm::cast<func::FuncOp>(m.lookupSymbol(cluster_func.getFunc()));
    for (auto entry : llvm::enumerate(cluster_func.getOperands())) {
      auto operand = SkipIdentityAndReadVariable(entry.value());
      auto block_arg = mlir::dyn_cast<BlockArgument>(operand);