std::next(copy_range.begin()) == copy_range.end() &&
      isa<tf_executor::GraphOp>(*copy_range.begin())) {
    // Already a graph.
    return;
  }

  auto return_op = dyn_cast<func::ReturnOp>(body.getTerminator());
  if (!return_op) {
    LLVM_DEBUG(llvm::dbgs() << "Expect function to end with return\n");
    return;
  }
  // Build GraphOp.
  OpBuilder builder(&body, body.begin());

    BlockArgument arg = std::get<1>(sharding_and_arg);
    if (failed(VerifySharding(arg.getType(), sharding))) return mlir::failure();
  }
  Operation* terminator = function_block.getTerminator();
  for (auto sharding_and_retval :
       llvm::zip(sharding_for_rets, terminator->getOpOperands())) {
    const auto& sharding = std::get<0>(sharding_and_retval);

void RewriteQuantizedIOPass::runOnOperation() {
  ModuleOp module = getOperation();
  OpBuilder builder(module);
  module.walk([&](func::FuncOp func) {
    Block& block = func.front();
    Operation* terminator = block.getTerminator();

    // Replace input_arg(tensor<quant_type>) -> tfr.cast
    // with input_arg(tensor<storage_type>) -> tfr.cast
    for (BlockArgument arg : block.getArguments()) {
      Type arg_type = arg.getType();

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<mlir::tf_device::ReturnOp>(loc, op->getResults());

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

  builder->restoreInsertionPoint(insert_point);

  return launch;
}

  FunctionType type;
  if (passthru_extra_args) {
    type = FunctionType::get(context, types, types);
  } else {
    SmallVector<Type, 4> result_types;
    auto operands = region.front().getTerminator()->getOperandTypes();
    result_types.append(operands.begin(), operands.end());
    type = FunctionType::get(context, types, result_types);
  }

  auto outlined_func = builder.create<func::FuncOp>(loc, name, type);

  op->replaceAllUsesWith(launch);

  builder->setInsertionPointToEnd(&launch.GetBody());
  builder->create<tf_device::ReturnOp>(loc, op->getResults());

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

llvm::SmallVector<Operation*, 4> GetResourceOpsUsingCompositeArgsInReplicate(
    tf_device::ReplicateOp replicate) {
  llvm::SmallVector<Operation*, 4> resource_users;

    unsigned region_index) {
  int num_region_results =
      GetRegionBlockWithIndex(region_index).getTerminator()->getNumOperands();

  int return_value_offset = 0;
  for (int region_id = 0; region_id < region_index; ++region_id)
    return_value_offset +=
        GetRegionBlockWithIndex(region_id).getTerminator()->getNumOperands();

  return getResults().slice(return_value_offset, num_region_results);
}

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

    llvm::ArrayRef<TF::AssignVariableOp> tail_assign_variable_ops) {
  const auto num_replicas = replicate_op.getN();
  auto return_op = llvm::dyn_cast<tf_device::ReturnOp>(
      replicate_op.getRegion().front().getTerminator());

  // Get the new result types.
  // TODO(prakalps): Do not add a value to returned values if it is already
  // returned.
  auto new_result_types = llvm::to_vector<4>(replicate_op->getResultTypes());