std::get<0>(result).replaceAllUsesWith(std::get<1>(result));
      }

      rewriter.setInsertionPoint(cloned_func.getBody().front().getTerminator());
      rewriter.create<TF::_XlaSendFromHostOp>(
          func.getLoc(),
          cloned_func.getBody().front().getTerminator()->getOperands(),
          /*dynamic_key=*/dynamic_key, op.getRecvKeyAttr(),
          /*device_ordinal=*/rewriter.getI64IntegerAttr(0),

  llvm::SmallVector<Type, 8> arg_types;
  for (auto arg : func.getArguments()) arg_types.push_back(arg.getType());
  func.setType(
      FunctionType::get(func.getContext(), arg_types,
                        func.front().getTerminator()->getOperandTypes()));
}

// Holds the size value of a tensor list and whether the size is statically
// known (fixed).
struct SizeInfo {
  Value size;
  bool fixed;
};

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

// Returns the aliasing argument number of a fucntion return value if it simply
// forwards the argument. Otherwise, returns -1.
int64_t FindAliasedInput(func::FuncOp func, int64_t return_index) {
  Value return_val = func.front().getTerminator()->getOperand(return_index);
  auto maybe_arg = mlir::dyn_cast<BlockArgument>(return_val);
  if (!maybe_arg) return -1;
  return maybe_arg.getArgNumber();
}

// 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;

void UpdateTerminatorArguments(T& func,
                               const ArrayRef<unsigned> arguments_to_erase,
                               llvm::BitVector& erase_indices) {
  auto terminator = func.front().getTerminator();
  int num_operands = terminator->getNumOperands();
  erase_indices.resize(num_operands);
  for (auto arg_index : arguments_to_erase) {
    auto argument = func.getArgument(arg_index);

  // 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;
}

  Block& block = region.front();
  if (std::distance(block.begin(), block.end()) != 2) return std::nullopt;
  TF::YieldOp yield =
      llvm::dyn_cast_or_null<TF::YieldOp>(block.getTerminator());
  if (!yield) return std::nullopt;
  func::CallOp call = llvm::dyn_cast_or_null<func::CallOp>(*block.begin());
  if (!call) return std::nullopt;
  if (block.getNumArguments() != call.getNumOperands() ||

  OpBuilder builder(replicate);
  auto new_replicate = builder.create<tf_device::ReplicateOp>(
      replicate.getLoc(), num_replicas, devices, new_replicated_inputs,
      new_packed_inputs,
      replicate.GetBody().getTerminator()->getOperandTypes());
  for (auto arg : replicate.GetBody().getArguments()) {
    if (replicate.IsReplicatedBlockArgument(arg)) {
      arg.replaceAllUsesWith(