ConversionPatternRewriter *rewriter) const {
    auto guard = OpBuilder::InsertionGuard(*rewriter);
    auto inputs = branch_func.getFunctionType().getInputs();
    Block *block = rewriter->createBlock(
        &branch_func.getBody(), branch_func.begin(), inputs,
        SmallVector<Location>(inputs.size(), branch_func.getLoc()));

  // Now use the filtered original operands, which will be replaced by
  // AddLoadsStoresOutsideControlFlowOp().
  auto new_while = builder.create<TF::WhileOp>(
      while_op.getLoc(), body.getFunctionType().getResults(),
      FilterRange<Value, OperandRange>(while_op.getOperands(),
                                       resource_arg_uses),
      while_op->getAttrs());
  // Prepare for AddLoadsStoresOutsideControlFlowOp().

  region.push_back(new mlir::Block());
  Location loc = region.getLoc();
  auto inputs = func.getFunctionType().getInputs();
  region.addArguments(inputs, mlir::SmallVector<Location>(inputs.size(), loc));
  op_builder.setInsertionPointToStart(&region.front());
  auto call_op = op_builder.create<mlir::func::CallOp>(
      loc, func.getFunctionType().getResults(), func.getSymName(),
      region.getArguments());

    // DCE pass.
    OpBuilder builder(module_);
    builder.setInsertionPointToEnd(&module_.getBodyRegion().back());
    const auto func_type =
        builder.getFunctionType(/*inputs=*/{}, /*results=*/{});
    auto summary_func = builder.create<func::FuncOp>(
        builder.getUnknownLoc(), /*sym_name=*/"summary", func_type);
    summary_func.setPrivate();

  llvm::SmallVector<Type, 4> output_types;
  output_types.reserve(outputs.size());
  for (Value v : outputs) output_types.emplace_back(v.getType());

  auto func_type = builder->getFunctionType(operand_types, output_types);

  FuncOp outlined_func =
      FuncOp::create(ops.front()->getLoc(), kHostFunctionAttr, func_type);

  // Create function body.

  auto out_types = GetValueTypes(outputs);
  builder.setInsertionPointToEnd(&module.getBodyRegion().back());
  auto func_op = builder.create<func::FuncOp>(
      module.getLoc(), name, builder.getFunctionType(in_types, out_types));
  func_op.setPrivate();
  symbol_table.insert(func_op);
  return func_op;
}

TF::StatefulPartitionedCallOp EncapsulateOpsInFunc(
    OpBuilder& builder, SymbolTable& symbol_table,