}

LogicalResult FindSourceTPUReplicatedOutput(
    Value val, TF::TPUReplicatedOutputOp& rep_out) {
  Operation* op = val.getDefiningOp();
  if (auto src = llvm::dyn_cast<TF::TPUReplicatedOutputOp>(op)) {
    rep_out = src;
    return LogicalResult::success();
  }
  if (auto src = llvm::dyn_cast<TF::IdentityOp>(op)) {
    return FindSourceTPUReplicatedOutput(src->getOperand(0), rep_out);
  }