for (auto branch : while_region.getRegions())
          add_to_worklist(branch->getArgument(use.getOperandNumber()));
        continue;
      }
      if (auto yield = llvm::dyn_cast<YieldOp>(use.getOwner())) {
        Operation* parent = yield->getParentOp();
        add_to_worklist(parent->getResult(use.getOperandNumber()));
        continue;
      }
      // TODO(jpienaar): This can be generalized.

  // This is v1 ref variable semantics in the TFLite runtime.
  bool is_variable = false;
  for (auto& use : value.getUses()) {
    is_variable = IsStatefulOperand(use.getOwner(), use.getOperandNumber());
    if (is_variable) {
      break;
    }
  }
  // The value is used as a variable if produced by an op with "tfl.is_variable"

    // Erase control outputs from feed.
    auto control_uses = island->getResult(1).getUses();
    for (auto& control_use : llvm::make_early_inc_range(control_uses))
      control_use.getOwner()->eraseOperand(control_use.getOperandNumber());

    if (!arg_node.node->requested_device().empty())
      arg_attrs[i].set("tf.device", builder_.getStringAttr(
                                        arg_node.node->requested_device()));