QuantizedType QuantizationDriver::GetQuantParamsForSameScaleConstraint(
    Operation* op) {
  // Two vector to collect Non-empty operands and results states.
  std::vector<QuantState*> mutable_states, immutable_states;
  for (int i = 0; i < op->getNumOperands(); ++i) {
    QuantState& state = GetOperandQuantState(op, i);
    if (state.immutable) {
      immutable_states.push_back(&state);
    } else if (!state.IsEmpty()) {

    // Returns the state of the index-th operand of the op.
    QuantState &GetOperandQuantState(Operation *op, int index) {
      return states_[operand_states_[{op, index}]];
    }

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

  }

  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:

  for (const auto& arg : quantization_driver.GetArgs()) {
    const QuantState& state = quantization_driver.GetArgQuantState(arg);
    EXPECT_TRUE(isa<quant::QuantizedType>(state.params));
  }
  for (const auto& result : quantization_driver.GetResultStates()) {
    Operation* op = result.first.first;
    const int res_index = result.first.second;
    const QuantState state =
        quantization_driver.GetResultQuantState(op, res_index);

// - use the first ready one in the collection.
QuantParams QuantizeContext::GetQuantParamsForSameScaleConstraint(
    Operation *op) {
  // Two vector to collect Non-empty operands and results states.
  std::vector<quant::QuantState *> mutable_states, immutable_states;
  for (int i = 0, e = op->getNumOperands(); i != e; ++i) {
    auto &state = states_manager_.GetOperandQuantState(op, i);
    if (state.immutable) {