}
  op->emitOpError() << "Value did not come from a TPUReplicatedOutput op: "
                    << val;
  return LogicalResult::failure();
}

int FindReturnIndex(Value val) {
  const int not_found = -1;
  for (auto user : val.getUsers()) {
    if (auto ret_op = llvm::dyn_cast<func::ReturnOp>(user)) {
      for (auto index = 0; index < ret_op->getNumOperands(); ++index) {