rewriter.replaceOp(bmm_op, sum_op);
    return success();
  };

 private:
  bool SplatValueEquals(SplatElementsAttr float_or_int, double rhs) const {
    if (mlir::isa<DenseFPElementsAttr>(float_or_int)) {
      return mlir::cast<DenseFPElementsAttr>(float_or_int)
          .getSplatValue<APFloat>()
          .isExactlyValue(rhs);
    } else if (mlir::cast<DenseIntElementsAttr>(float_or_int)) {

// Extracts min and max values from the DenseFPElementsAttr, and stores them
// into `mins` and `maxs`. When mins and maxs are extracted per-channel,
// `dim_size` is number of channels and `slice_size` is the size of slice per
// each channel. When `symmetric` is true, the range is expanded to [-M, M].
void ExtractMinMaxFromAttr(DenseFPElementsAttr values, int dim_size,

    QuantizedType quant_type;
    DenseFPElementsAttr attr;
    if (!matchPattern(op->getResult(0), m_Constant(&attr))) return false;

    if (attr.size() < quant_specs_.minimum_elements_for_weights) {
      op->emitRemark("Quantization is skipped for ")
          << quantized_op->getName().getStringRef().str() << " because it has "
          << mlir::dyn_cast<DenseFPElementsAttr>(attr).size()

// to `raw_value`.
template <typename T>
bool IsConstantValueOf(Value value, T raw_value) {
  auto element_type = mlir::cast<ShapedType>(value.getType()).getElementType();
  if (mlir::isa<FloatType>(element_type)) {
    DenseFPElementsAttr float_attr;
    if (matchPattern(value, m_Constant(&float_attr)) && float_attr.isSplat() &&
        float_attr.getSplatValue<APFloat>().isExactlyValue(raw_value))
      return true;

    }
  }
  return res;
}

// Gets a matrix which corrects the overcounting of divisors when casting a
// average pool with ceil mode true as one with ceil mode false on a padded
// tensor.
DenseFPElementsAttr GetCorrectionMatrix(Builder& builder, CompositeOp op) {
  const TorchAvgPoolData pool = GetTorchAvgPoolData(op);

  llvm::SmallVector<int64_t, 4> nhwc_out_shape(4);
  nhwc_out_shape[0] = 1;  // Broadcast batch.

    return builder.create<TF::ConstOp>(loc, attr);
  }

  const auto type = RankedTensorType::get(shape, builder.getF32Type());
  const auto value_attr = DenseFPElementsAttr::get(type, values);
  return builder.create<TF::ConstOp>(loc, value_attr);
}

// Creates a 1D array with integer/float type.
template <typename T>
Value Create1DConstValue(OpBuilder& builder, const Location loc,

      op.getOutput().replaceAllUsesWith(op.getInput());
      rewriter.eraseOp(op);
      return success();
    }

    // The layer stats contain only the first min/max pairs.
    ElementsAttr layer_stats = DenseFPElementsAttr::get(
        RankedTensorType::get({2}, rewriter.getF32Type()),
        {static_cast<float>(min.getValueAsDouble()),
         static_cast<float>(max.getValueAsDouble())});
    ElementsAttr axis_stats;

    max_value = tf_op.getMaxAttr();
    return true;  // Successfully matched and fetched.
  }
};

template <class TFFakeQuantOp>
struct FetchConstantMinMaxInputs {
  using AttrType = DenseFPElementsAttr;
  bool operator()(TFFakeQuantOp tf_op, AttrType &min_value,
                  AttrType &max_value) const {
    Value min = tf_op.getMin(), max = tf_op.getMax();
    if (!matchPattern(min, m_Constant(&min_value))) {

    // a special case that there are tf.Identity ops between the min/max
    // constants and the tf.FakeQuantWithMinMaxVarsOp.
    Value min = tf_op.getMin(), max = tf_op.getMax();
    DenseFPElementsAttr min_value, max_value;
    if (auto id1 = dyn_cast_or_null<TF::IdentityOp>(min.getDefiningOp())) {
      id1.replaceAllUsesWith(id1.getInput());
      min = tf_op.getMin();
      rewriter.eraseOp(id1);
    }

  for (int i = 0; i < op->getNumResults(); ++i) {
    if (GetResultQuantState(op, i).IsEmpty()) return false;
  }
  return true;
}

bool QuantizationDriver::SetConstantResultParams(Operation* op) {
  DenseFPElementsAttr attr;
  const Value result = op->getResult(0);
  if (!matchPattern(result, m_Constant(&attr))) {
    return false;
  }
  // TODO: b/323478683 - Make storage_type_width and narrow_range configurable.