Operation* tpu_copy_with_dynamic_shape_op) {
  llvm::SmallSetVector<Value, 4> new_launch_op_results;

  new_launch_op_results.insert(
      old_launch_op->GetBody().getTerminator()->getOperands().begin(),
      old_launch_op->GetBody().getTerminator()->getOperands().end());

  for (Value operand : tpu_copy_with_dynamic_shape_op->getOperands()) {
    if (GetOpOfValue(operand)->getParentRegion() ==

          GetOpOfValue(cluster_func_operand.value()));
      if (!device_launch_op) continue;
      for (auto result : llvm::zip(
               device_launch_op.getResults(),
               device_launch_op.GetBody().getTerminator()->getOperands())) {
        if (std::get<0>(result) == cluster_func_operand.value() &&
            llvm::isa<TF::TPUAnnotateTensorsWithDynamicShapeOp>(
                std::get<1>(result).getDefiningOp())) {

    if (!func_op) return failure();
    // The quantized fusion should have requantize and return ops at the end.
    auto return_op = dyn_cast_or_null<func::ReturnOp>(
        func_op.getRegion().getBlocks().front().getTerminator());
    if (!return_op) return failure();
    auto req_op = llvm::dyn_cast_or_null<mlir::stablehlo::UniformQuantizeOp>(
        return_op.getOperands()[0].getDefiningOp());
    if (!req_op) return failure();

                                            tf_device::LaunchOp launch) {
  // Forward launch inner op results to launch op results.
  launch.replaceAllUsesWith(launch.GetBody().getTerminator()->getOperands());

  // For all inner ops, assign the launch device as a `device` attribute.
  if (failed(AssignDevicesInRegion(tf_dialect, launch, launch.getBody())))
    return failure();

// op operands' types.
//
// Requires the function has exactly one block.
void UpdateFuncType(func::FuncOp func) {
  Operation* terminator = func.front().getTerminator();
  auto return_types = llvm::to_vector<4>(terminator->getOperandTypes());

  FunctionType func_type = func.getFunctionType();
  if (llvm::ArrayRef(return_types) == func_type.getResults()) return;

  // Replace `tf_device.return` terminator with `std.return` in function
  // body.
  auto return_op =
      cast<tf_device::ReturnOp>(outlined_func_block->getTerminator());
  builder->setInsertionPoint(return_op);
  ReplaceClusterReturnWithReturn(return_op, builder);

  symbol_table->insert(outlined_func);
  return outlined_func;
}

}

LogicalResult ModifyIONodesPass::ModifyOutputNodes(
    func::FuncOp func, llvm::SmallVectorImpl<Type>& new_output_types,
    OpBuilder builder) {
  Block& block = func.front();
  auto* terminator = block.getTerminator();
  builder.setInsertionPoint(terminator);

  if (mlir::isa<FloatType>(output_type)) {
    return success();
  }

  int num_return_operands = terminator->getNumOperands();