new_args.reserve(extern_values.size());
  Block& block = func_region.front();
  for (Value value : extern_values) {
    auto arg = block.addArgument(value.getType(), loc);
    replaceAllUsesInRegionWith(value, arg, func_region);
    new_args.push_back(arg);
  }

  // Replace yield op with return.
  Operation* yield_op = outlined_func.getBody().front().getTerminator();
  OpBuilder b(yield_op);

  // Replace uses of live-in values within cluster_op region with function
  // arguments.
  Region& op_region = op.getBody();
  for (auto p : llvm::zip(live_ins, outlined_func_block->getArguments())) {
    replaceAllUsesInRegionWith(std::get<0>(p), std::get<1>(p), op_region);
  }

  // Move all instructions in cluster_op into outlined_function's only block.
  auto& op_body = op.GetBody().getOperations();

      auto loc = outlined_func.getLoc();
      for (Value operand : operands) {
        BlockArgument newArg = entry_block.addArgument(operand.getType(), loc);
        replaceAllUsesInRegionWith(operand, newArg, outlined_func.getBody());
      }
    }

    // The function is in place in the nested module, create a call and yield in
    // the original island.

      for (OpResult result : it->getResults()) {
        replaceAllUsesInRegionWith(
            result,
            it->getThenBranch().front().getTerminator()->getOperand(
                result.getResultNumber()),
            it2->getThenBranch());
        replaceAllUsesInRegionWith(
            result,
            it->getElseBranch().front().getTerminator()->getOperand(

      }
      // Add constant at start of region.
      auto const_builder = OpBuilder::atBlockBegin(&region->front());
      auto const_value = const_builder.clone(*extern_value.getDefiningOp());
      replaceAllUsesInRegionWith(extern_value, const_value->getResult(0),
                                 *region);
    }
  }

  return llvm::to_vector<4>(extern_values);
}

      // If we found matching input ops, we can remove the TPUReplicatedInput
      // ops and replace their result values with the inputs to the matching
      // TPUReplicatedOutput op.
      replaceAllUsesInRegionWith(input_op.getResult(), output_op.getOperand(),
                                 func.getRegion());
      ops_to_erase.insert(input_op);

      // If the outputs aren't also returned from this function, then we can

  for (Operation* op : topological_order) op->moveBefore(block, block->end());

  // Replace the 'inputs' values with the new function's arguments.
  for (auto p : llvm::zip(inputs, new_func.getArguments()))
    replaceAllUsesInRegionWith(std::get<0>(p), std::get<1>(p),
                               new_func.getBody());

  builder.setInsertionPointToEnd(block);
  builder.create<func::ReturnOp>(parent_func.getLoc(), outputs.getArrayRef());

      builder, cluster, /*before=*/true, launch_block, host_device);

  for (auto result : llvm::zip(launch.GetBody().getTerminator()->getOperands(),
                               launch.getResults()))
    replaceAllUsesInRegionWith(std::get<0>(result), std::get<1>(result),
                               cluster.getBody());
}

// Extracts and move outside compiled ops that have no dependencies in the
// cluster to before the cluster.

                 replicate_op.GetBody().getArguments())) {
    mlir::TF::TPUReplicatedInputOp input = std::get<0>(input_and_block_arg);
    Value block_arg = std::get<1>(input_and_block_arg);
    mlir::replaceAllUsesInRegionWith(input->getResult(0), block_arg,
                                     cluster.getBody());
    // Update replicated input use in tf.TPUPartitionedInputV2 op.
    for (auto& use : input->getUses()) {
      auto pi =