mlir::RegisterCommonToolingDialects(registry_);
    context_.appendDialectRegistry(registry_);
    context_.loadAllAvailableDialects();
  }

  absl::StatusOr<std::vector<mlir::tf_device::ClusterFuncOp>> GetClusterFuncOps(
      absl::string_view mlir_module_filename) {
    TF_RETURN_IF_ERROR(CreateMlirModule(mlir_module_filename));
    std::vector<mlir::tf_device::ClusterFuncOp> cluster_func_ops;