}

    def "prefers to stop less frequently used idle clients when releasing memory"() {
        def client1 = Mock(WorkerDaemonClient) { _ * getUses() >> 5 }
        def client2 = Mock(WorkerDaemonClient) { _ * getUses() >> 1 }
        def client3 = Mock(WorkerDaemonClient) { _ * getUses() >> 3 }
        starter.startDaemon(options) >>> [client1, client2, client3]

    llvm::MapVector<Operation*, llvm::SmallVector<unsigned int, 4>>*
        arguments_to_erase) {
  auto arg = region->getArgument(argument_index);
  for (auto& usage : arg.getUses()) {
    auto* user_op = usage.getOwner();
    int operand_index = usage.getOperandNumber();
    PropagateUsage(user_op, operand_index, value, work_list,
                   arguments_to_erase);
  }
}

    const ShapedType type = mlir::dyn_cast<ShapedType>(op.getType());
    if (!type || !type.getElementType().isF32()) return quantizable_ops;

    const Value value = op.getResult();

    for (OpOperand& use : value.getUses()) {
      Operation* user = use.getOwner();
      const int operand_num = use.getOperandNumber();
      quantizable_ops.insert({user, operand_num});
    }
    return quantizable_ops;
  }

  for (auto& use : uses_optional.value()) {
    if (llvm::isa<TF::WhileOp>(use.getUser())) {
      if (*while_op != nullptr) {
        use.getUser()->emitOpError() << "multiple users of function.";
        return LogicalResult::failure();
      } else {
        *while_op = llvm::cast<TF::WhileOp>(use.getUser());
      }
    } else {
      use.getUser()->emitOpError() << "non while use of function.";

// -------------------------------------------------------------------------- //

namespace {
constexpr char kDeviceAttr[] = "device";

// A type that abstracts over types that have uses accessible via `getUses`.
using Source = PointerUnion<Operation *, BlockArgument *>;

// We use union-find algorithm to build clusters of connected operations based
// on the user provided policy. If an operation can be clustered (one of the

// multiple such activations, one is returned (with no guarantee as to which
// one). If there are no activation functions that use the output, returns
// nullptr.
Operation *GetActivation(Value op) {
  for (auto &use : op.getUses()) {
    if (IsActivationFunction(use.getOwner())) return use.getOwner();
  }
  return nullptr;
}

// Finds and returns a BiasAdd that uses the result of `op` as the `value`

    ShapedType new_result_type = old_result_type.clone(quantized_type);

    // Insert CastOp if it does not exist yet. Otherwise, just rewire without
    // creating a CastOp.
    for (auto& connected_op : op.getResult().getUses()) {
      auto cast_op = llvm::dyn_cast_or_null<CastOp>(connected_op.getOwner());
      if (cast_op && cast_op.getType() == new_result_type) {
        quantize_op->setOperand(quantize_operand_num, cast_op);
        return;

  // terminator.
  bool checkIfAnyUserIsConnectedToTermiantor(BlockArgument op) const {
    for (const auto& user : op.getUsers()) {
      if (user->template hasTrait<OpTrait::IsTerminator>()) return true;
      if (auto next_user = dyn_cast_or_null<TF::IdentityOp>(user)) {
        return (*(next_user->getResult(0).getUsers().begin()))
            ->template hasTrait<OpTrait::IsTerminator>();
      }
    }
    return false;
  }

  xla::OpSharding sharding;
  sharding.ParseFromString(
      old_partitioned_input.get_XlaShardingAttr().getValue().str());
  for (OpOperand& read_use :
       llvm::make_early_inc_range(old_read.getValue().getUses())) {
    if (dyn_cast_or_null<tf_device::ClusterFuncOp>(read_use.getOwner())) {
      // ClusterFunc's use of the Read is replaced with use of the
      // TPUPartitionedInputV2.
      read_use.set(new_partitioned_input);

        synchronized (lock) {
            List<WorkerDaemonClient> sortedClients = CollectionUtils.sort(idleClients, comparingInt(WorkerDaemonClient::getUses));
            List<WorkerDaemonClient> clientsToStop = selectionFunction.transform(new ArrayList<>(sortedClients));
            if (!clientsToStop.isEmpty()) {
                stopWorkers(clientsToStop);
            }