}

  return nullptr;
}

/// Converts a real expressed DenseFPElementsAttr to a corresponding
/// DenseElementsAttr (typically DenseIntElementsAttr) containing quantized
/// storage values assuming the given quantizedElementType and converter.
static DenseElementsAttr convertDenseFPElementsAttr(
    DenseFPElementsAttr realFPElementsAttr,
    quant::QuantizedType quantizedElementType,

}

ElementsAttr UniformQuantizedPerAxisValueConverter::convert(
    Attribute real_value) {
  if (auto attr = dyn_cast<DenseFPElementsAttr>(real_value)) {
    return convert(attr);
  }
  return nullptr;
}

DenseElementsAttr UniformQuantizedPerAxisValueConverter::convert(
    DenseFPElementsAttr attr) {
  // Creates the converter for each chunk. Normally the size of the

    return success(
        mlir::isa_and_nonnull<DenseFPElementsAttr>(const_op.getValue()));
  }

  void rewrite(mlir::stablehlo::TransposeOp op,
               PatternRewriter& rewriter) const override {
    auto const_op =
        cast<mlir::stablehlo::ConstantOp>(op.getOperand().getDefiningOp());

    const auto value_attr =
        mlir::cast<DenseFPElementsAttr>(const_op.getValue());

// Get dense attr for a matrix that corrects the over counting of divisors when
// casting an average pool with ceil mode on in terms of average pool with it
// off.
DenseFPElementsAttr GetCorrectionMatrix(Builder& builder, mhlo::CompositeOp op);

}  // namespace odml
}  // namespace mlir

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

// 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.

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

  LogicalResult matchAndRewrite(mlir::stablehlo::UniformQuantizeOp op,
                                PatternRewriter& rewriter) const override {
    DenseFPElementsAttr attr;
    if (matchPattern(op.getOperand(), m_Constant(&attr))) {
      const Type qtype = op.getResult().getType();
      ElementsAttr quantized_attr = Quantize(attr, qtype);
      if (quantized_attr) {