mlir::OwningOpRef<mlir::ModuleOp> module =
      mlir::parseSourceString<mlir::ModuleOp>(code, &context);
  ASSERT_TRUE(module);
  EXPECT_TRUE(HasTPUPartitionedCallOpInModule(*module));
}

TEST(HasTPUPartitionedCallOpInModule, HasNotTPUPartitionedCallModule) {
  const char* const code = R"mlir(
module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

bool IsSupportedByReplicatedBridge(mlir::ModuleOp module);

// Check if an MLIR module contains TPUPartitionedCall op. If so, we define
// such graph as an inference graph. Otherwise, it is non inference graph.
bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module);

// Check if a graph contains TPUPartitionedCall op, including its reachable
// functions. The function library is used to store the functions that are
// defined in a TensorFlow program

  return IsReplicatedGraph(graph, function_library);
}

bool IsSupportedByReplicatedBridge(mlir::ModuleOp module) {
  return IsReplicatedGraph(module);
}

bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module) {
  bool has_tpu_partitioned_call = false;
  for (auto func_op : module.getOps<mlir::func::FuncOp>()) {
    func_op->walk([&](mlir::TF::TPUPartitionedCallOp op) {