TF::TPUReplicatedInputOp private_input = input.clone();
    builder.insert(private_input);
    forward_pass_ops.insert(private_input);
    for (OpOperand& next_use : input.getOutput().getUses()) {
      if (!forward_pass_ops.contains(next_use.getOwner())) continue;
      next_use.getOwner()->setOperand(next_use.getOperandNumber(),
                                      private_input.getOutput());
    }
  }

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

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

    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;

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

    // determine whether we have taken care of all uses and can remove
    // the arg.
    for (BlockArgument arg : func.getArguments()) {
      auto uses = arg.getUses();
      use_count[arg.getArgNumber()] = std::distance(uses.begin(), uses.end());
      argument_to_index.insert({arg, arg.getArgNumber()});
    }

    // We're only considering return values that are the same (at this

  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);

  for (auto ret_vals : llvm::zip(results, cluster.getResults())) {
    Value old_ret = std::get<0>(ret_vals);
    Value new_ret = std::get<1>(ret_vals);
    for (auto& use : llvm::make_early_inc_range(old_ret.getUses())) {
      Operation* user = use.getOwner();
      if (!body->findAncestorOpInBlock(*user)) use.set(new_ret);
    }
  }

  // Move ops that depend on something in the cluster behind the cluster.

	return kvh.kubelet.nodeName
}

func (kvh *kubeletVolumeHost) GetEventRecorder() record.EventRecorder {
	return kvh.kubelet.recorder
}

func (kvh *kubeletVolumeHost) GetExec(pluginName string) utilexec.Interface {
	return kvh.exec

    while (!uses_at_current_level.empty()) {
      llvm::SmallVector<mlir::Value> next_values_to_visit;
      for (auto cur_op : uses_at_current_level) {
        for (auto& cur_op_use : cur_op.getUses()) {
          Operation* next_op = cur_op_use.getOwner();
          int next_op_operand_num = cur_op_use.getOperandNumber();
          if (auto call_op = llvm::dyn_cast<mlir::CallOpInterface>(next_op)) {

// if there is only one.
Operation* GetUserIfOnlyOne(Operation* op) {
  if (op->getNumResults() != 1) return nullptr;
  auto result = op->getResult(0);
  if (!result.hasOneUse()) return nullptr;
  return (*result.getUses().begin()).getOwner();
}

// Gets operation providing value for the given operand of given operation
// if the given operation is the only user.
Operation* GetInputOpWithOneUse(Operation* op, int opr_num) {