std::vector<Value>& n_operands) {
  n_operands.clear();
  n_operands.reserve(non_tpu_caller->getNumOperands());
  for (auto operand : non_tpu_caller->getOperands()) {
    if (llvm::isa<BlockArgument>(operand)) {
      auto arg = llvm::cast<BlockArgument>(operand);
      n_operands.push_back(loop_operands[arg.getArgNumber()]);
      continue;
    }

  for (const auto& pos_and_input : llvm::enumerate(replicated_input_ops)) {
    auto input = pos_and_input.value();
    bool is_packed = input.getIsPacked();
    const int num_operands = input->getNumOperands();
    int num_inputs = is_packed ? 1 : num_replicas;
    if (num_operands != num_inputs)
      return input->emitOpError() << "requires " << num_inputs << " operands";
    if (is_packed) {
      packed_inputs.push_back(input->getOperand(0));

  // Ensure that the yield terminator operands matches the island results type.
  int result_count = island.getNumResults() - 1;  // -1 for the control token
  const int num_operands = yield.getNumOperands();
  if (num_operands != result_count)
    return yield.emitOpError()
           << "has " << yield.getNumOperands()
           << " operand, but island returns " << result_count;

  // Adds the "source" to the operands of the dst by creating a new dst
  // operation.
  mlir::OperationState state(dst->getLoc(), dst->getName());
  auto num_operands = dst->getNumOperands();
  state.operands.reserve(num_operands + 1);
  for (int input = 0, e = num_operands + 1; input != e; ++input) {
    if (input < dst_input) {
      state.operands.push_back(dst->getOperand(input));
    } else if (input == dst_input) {