py::overload_cast<mlir::Block*, mlir::Block::iterator>(
               &mlir::OpBuilder::setInsertionPoint))
      .def("saveInsertionPoint", &mlir::OpBuilder::saveInsertionPoint)
      .def("restoreInsertionPoint", &mlir::OpBuilder::restoreInsertionPoint)
      .def(
          "create",
          [](mlir::OpBuilder& opb, mlir::OperationState& state) {
            return opb.create(state);
          },

  builder->create<mlir::tf_device::ReturnOp>(loc, op->getResults());

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

  builder->restoreInsertionPoint(insert_point);

  return launch;
}

    def __init__(self, arg0: Block, arg1: Block_Iterator) -> None: ...
    def create(self, *args, **kwargs) -> Any: ...
    def getContext(self) -> MLIRContext: ...
    def getUnknownLoc(self) -> Location: ...
    def restoreInsertionPoint(self, arg0) -> None: ...
    def saveInsertionPoint(self, *args, **kwargs) -> Any: ...
    def setInsertionPoint(self, arg0: Block, arg1: Block_Iterator) -> None: ...

class OpBuilder_InsertionPoint:

  symbol_table.insert(quantization_func);

  FlatSymbolRefAttr func_name_attr =
      FlatSymbolRefAttr::get(rewriter.getStringAttr(func_name));

  rewriter.restoreInsertionPoint(original_point);

  auto quantize_call = rewriter.create<TF::PartitionedCallOp>(
      quantized_op->getLoc(), quantize_result_type, input, func_name_attr,

  std::vector<Operation*> topological_order;
  for (Operation& op : parent_func.getOps())
    if (ops.contains(&op)) topological_order.push_back(&op);

  // Create the partitioned call
  builder.restoreInsertionPoint(saved_insertion_point);
  auto caller = MakeFuncCaller(builder, module.getLoc(), new_func,
                               inputs.getArrayRef(), flag_for_inlining);

  Block* block = new_func.addEntryBlock();

  builder->create<tf_device::ReturnOp>(loc, op->getResults());

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

  builder->restoreInsertionPoint(insert_point);
}

// Performs the transformation for a replicate op inside a while loop. Returns
// true when any change was made by this function.
bool HandleReplicateOp(TF::WhileRegionOp while_op,

  std::vector<Operation*> topological_order;
  for (Operation& op : parent_func.getOps())
    if (ops.contains(&op)) topological_order.push_back(&op);

  // Create the partitioned call
  builder.restoreInsertionPoint(insertion_point);
  auto caller = MakeFuncCaller(builder, module.getLoc(), new_func, inputs);

  Block* block = new_func.addEntryBlock();

    // Using SymbolTable will ensure that the inserted symbol names are
    // unique.
    manager.insert(then_branch_op);
    manager.insert(else_branch_op);

    rewriter.restoreInsertionPoint(original_point);
    rewriter.replaceOpWithNewOp<TF::IfOp>(
        op, result_type, if_cond,
        /*input=*/
        ArrayRef<Value>({input_handle, input_shape, size_diff, size}),
        /*then_branch=*/

  mlir::Operation* cloned_op = op_builder.clone(*op);
  // Add the yield operation.
  op_builder.create<mlir::TFL::YieldOp>(op_loc, cloned_op->getResults());
  // Now emit into the function body again.
  op_builder.restoreInsertionPoint(saved_pos);

  // The ControlNodeOp depends on all control tokens emitted by the nodes at the
  // other end of the incoming edges. Since we're proceding in a valid