/*config=*/builder.getStringAttr(""),
        /*config_proto=*/builder.getStringAttr(""),
        /*executor_type=*/builder.getStringAttr(""));
    SmallVector<Value, 16> yield_operands(call_op.getResults());
    builder.create<YieldOp>(island_op.getLoc(), yield_operands);
  }

  // Outline all the transitively called functions by moving them in the
  // outlined module.
  for (func::FuncOp func : outlined_module.getOps<func::FuncOp>()) {

// Creates respective YieldOp for the new merged island.
YieldOp CreateNewIslandYieldOp(IslandOp new_island,
                               llvm::MutableArrayRef<IslandResult> results) {
  llvm::SmallVector<Value, 8> yield_operands;
  yield_operands.reserve(results.size());

  for (auto [old_result, new_island] :
       llvm::zip(results, new_island.getOutputs())) {
    // Replace original island result with new island result.

            return !wrapped_ops.count(user.getOwner());
          }))
        yield_operands.push_back(std::get<1>(result));
    }
  }
  OpBuilder::atBlockEnd(&island_body)
      .create<YieldOp>(new_island.getLoc(), yield_operands);

  // remap results of the new islands to the user outside of the island.
  int current_result = 0;

      int result_number = llvm::cast<OpResult>(operand).getResultNumber();
      if (auto pred_island = llvm::dyn_cast<tf_executor::IslandOp>(pred)) {
        Value yield_operand = pred_island.GetYield().getOperand(result_number);
        predecessors.push_back(yield_operand.getDefiningOp());
      }
    }
  }
  return predecessors;
}

struct TPUColocateSplits
    : public impl::TPUColocateSplitsPassBase<TPUColocateSplits> {

    Operation *yield_op = body.front().getTerminator();
    Value yield_operand = yield_op->getOperand(result_idx);
    Value while_operand = op.getOperand(result_idx);
    Value body_arg = body.getArgument(result_idx);
    Value cond_arg = cond.getArgument(result_idx);
    if (yield_operand != body_arg && yield_operand != while_operand) {
      return op.emitOpError("Result #") << result_idx << " is not tied to arg #"

    // yield %cond, [...passthrough args...]
    if (yield->getNumOperands() != block.getNumArguments() + 1)
      return std::nullopt;
    for (auto [yield_operand, block_arg] :
         llvm::zip(yield->getOperands().drop_front(1), block.getArguments())) {
      if (yield_operand != block_arg) return std::nullopt;
    }
  } else {
    // All call results should feed into expected consumer