OpBuilder builder(&body, body.begin());
  auto graph_op = builder.create<tf_executor::GraphOp>(
      loc, func.getFunctionType().getResults());
  graph_op.getBody().push_back(new Block);
  builder.setInsertionPointToEnd(&graph_op.GetBody());
  auto island = builder.create<tf_executor::IslandOp>(
      loc, func.getFunctionType().getResults(),
      tf_executor::ControlType::get(&getContext()), ArrayRef<Value>());
  // Create Fetch.

  auto parallel_execute = builder->create<tf_device::ParallelExecuteOp>(
      cluster_func.getLoc(), 1, output_types);
  cluster_func->remove();
  auto& block = parallel_execute.GetRegionBlockWithIndex(0);
  builder->setInsertionPointToEnd(&block);
  builder->insert(cluster_func);
  cluster_func.replaceAllUsesWith(parallel_execute);
  builder->create<tf_device::ReturnOp>(block.getParent()->getLoc(),

  // condition.
  builder.setInsertionPointToEnd(orig_block_head);
  JumpToBlock(loc, get_operand, cond_block, &builder);

  // Call condition function in the condition block and then branch to the body
  // block or remainder of the original block depending on condition function
  // result.
  builder.setInsertionPointToEnd(cond_block);

  StringRef dialect = ops[0]->getName().getDialectNamespace();
  OpBuilder builder(context);
  // Every ModuleOp has at least one region and one block.
  Block* first_block = &module->getRegion(0).front();
  builder.setInsertionPointToEnd(first_block);
  auto func = builder.create<mlir::func::FuncOp>(
      ops[0]->getLoc(), dialect.str() + std::to_string(function_id),
      builder.getFunctionType(input_types, output_types));

  then_branch.push_back(new Block);
  builder.setInsertionPointToEnd(&then_branch.front());
  builder.create<mlir::TF::YieldOp>(if_region.getLoc(),
                                    /*operands=*/ArrayRef<Value>{});

  // Create empty else branch region.
  auto& else_branch = host_side_if.getElseBranch();
  else_branch.push_back(new Block);
  builder.setInsertionPointToEnd(&else_branch.front());

                                executor_graph.getBody().end());

    // Add terminator of tf_executor::graph
    rewriter.setInsertionPointToEnd(&executor_graph.getBody().front());
    rewriter.create<tf_executor::FetchOp>(loc);

    // Add terminator of func
    rewriter.setInsertionPointToEnd(&func.getBody().front());
    rewriter.create<func::ReturnOp>(loc);

  // Empty else branch, if there is no more data, do nothing.
  auto& else_branch = dataset_if.getElseBranch();
  else_branch.push_back(new Block);
  builder.setInsertionPointToEnd(&else_branch.front());
  // Return only the state variables from the body arguments.
  SmallVector<Value, 4> else_returns;
  for (int i = 1; i < state_size + 1; i++) {
    else_returns.push_back(body_args[i]);

      op->getLoc(), builder.getStringAttr(device),
      /*result_types=*/op->getResultTypes());
  op->replaceAllUsesWith(launch_op);

  launch_op.getBody().push_back(new Block);
  builder.setInsertionPointToEnd(&launch_op.GetBody());
  auto* return_op =
      builder.create<tf_device::ReturnOp>(op->getLoc(), op->getResults())
          .getOperation();
  MLIRContext* context = launch_op.getContext();

  auto launch = builder->create<tf_device::LaunchOp>(
      loc, builder->getStringAttr(device), op->getResultTypes());
  launch.getBody().push_back(new Block);
  op->replaceAllUsesWith(launch);

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<tf_device::ReturnOp>(loc, op->getResults());

  // Move op inside cluster.
  op->moveBefore(launch.GetBody().getTerminator());
}

  OpBuilder builder(module.getContext());

  OwningOpRef<ModuleOp> pruned =
      builder.create<ModuleOp>(module->getLoc());
  (*pruned)->setAttrs(module->getAttrs());
  builder.setInsertionPointToEnd(pruned->getBody());

  llvm::SmallDenseSet<func::FuncOp> added;
  for (const llvm::StringRef function_name : function_names) {
    auto func =