};

// Prunes result defining op if possible, returns true if pruning was done.
bool TryPruneResultDefiningOp(TF::WhileRegionOp while_op, OpResult result) {
  // Don't prune if result is used.
  if (!result.use_empty()) return false;

  Block& body_block = while_op.getBody().front();
  Block& cond_block = while_op.getCond().front();
  Operation* body_yield_op = body_block.getTerminator();

      for (mlir::Operation* user :
           llvm::make_early_inc_range(result.getUsers())) {
        if (llvm::isa<mlir::TF::TPUPartitionedOutputV2Op>(user)) {
          assert(user->use_empty());
          user->erase();
        }
      }
    }

    for (auto operand : cluster.getOperands()) {
      mlir::Operation* def = operand.getDefiningOp();
      if (operand.hasOneUse() &&

// results of the operation will be eliminated.
void EliminateUnusedResults(
    Operation *op, const llvm::BitVector *results_to_eliminate = nullptr) {
  auto can_eliminate = [&](OpResult &result) -> bool {
    if (!result.use_empty()) return false;
    if (results_to_eliminate)
      return results_to_eliminate->test(result.getResultNumber());
    else
      return true;
  };
  SmallVector<Type, 4> new_result_types;

  // i is consumed element if result(i) is used outside whileOp or
  // argument(i) is used in whileOp.getCond().
  for (auto i = 0; i < n; ++i) {
    if (!while_op.getResult(i).use_empty() ||
        !cond.getArgument(i).use_empty()) {
      explicitly_consumed_ids.push_back(i);
    }
  }
  // Empty consumed_element_ids implies none of results is used.
  if (explicitly_consumed_ids.empty()) {
    while_op.erase();

        assert(replicated_input->use_empty());
        replicated_input->erase();
        return WalkResult::advance();
      });

  if (result.wasInterrupted()) {
    signalPassFailure();
    return;
  }

  getOperation()->walk([](TF::TPUPartitionedInputV2Op partitioned_input) {
    if (partitioned_input->use_empty()) partitioned_input->erase();
  });
}

}  // namespace

        quantize_op.emitError(" Couldn't be modified to the requested type.");
        return failure();
      }
      new_input_types[i] = arg_type;
      arg.dropAllUses();
      if (quantize_op.use_empty()) {
        quantize_op.erase();
      }
    } else {
      // `arg` has multiple uses or the user isn't a quantiz op (so we couldn't
      // rewrite it to a different type. Make a copy of the `arg` and replace

      op->getLoc(), op->getResultTypes()[0], name, /*container=*/"",
      op.getOutputTypes(), op.getOutputShapes());
  op->getResults()[0].replaceAllUsesWith(new_op->getResults()[0]);
  if (op->use_empty()) op->erase();
  return count;
}

void NameAnonymousIteratorsPass::runOnOperation() {
  int count = 1;
  getOperation().walk(
      [&](TF::AnonymousIteratorOp op) { count = replace(op, count); });

        auto new_op = builder.create<TFR::CastOp>(
            returned_op->getLoc(), new_type, returned_op.getArg());
        returned_value.set(new_op.getResult());
        if (returned_op.use_empty()) {
          returned_op.erase();
        }
      } else {
        returned_value.get().getDefiningOp()->emitError(
            "The producer of quantized type result should be a tfr.cast op.");
        return;

      status = ReplacePartitionedOp(num_cores_per_replica, partitioned_output);
    }

    if (status.has_value()) {
      if (failed(*status) || !op->use_empty()) return WalkResult::interrupt();

      op->erase();
    }

    return WalkResult::advance();
  });

  if (result.wasInterrupted()) {
    signalPassFailure();
    return;
  }
}

  // Add sink island to pin all islands as a control dependency if there is a
  // control dependency leading from the parallel_execute originally.
  if (!island_op.getControl().use_empty()) {
    llvm::SmallVector<Value, 8> island_operands;
    for (auto& execute : executes)
      island_operands.push_back(execute.getControl());

    builder.setInsertionPoint(island_op);