}

  DenseMap<OpWithResultIndex, QuantStateIndex> result_states_;

  // Returns the state of the block argument.
  QuantState& GetArgQuantState(BlockArgument arg) {
    return states_[arg_states_[arg]];
  }

  // Returns the state of the index-th result of the op.
  QuantState& GetResultQuantState(Operation* op, const int index) {
    return states_[result_states_[{op, index}]];
  }

 private:

  if (!inserted) {
    arg_states_[arg] = cached->second;
    return;
  }

  const QuantizedType quantized_type =
      QuantizedType::getQuantizedElementType(arg_value.getType());
  const bool immutable = quantized_type != nullptr;
  const QuantizationDriver::QuantStateIndex next_state_index = states_.size();
  states_.push_back({quantized_type, immutable});
  arg_states_[arg] = next_state_index;