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,

  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();