// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-quantize | FileCheck %s

// Test that hybrid quantized dot_general is produced when q/dq pair only exists
// for weight.

module attributes {tf_saved_model.semantics} {
  func.func private @quantize_dot_general_fn(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {
    %cst = stablehlo.constant dense<3.000000e-01> : tensor<2x3xf32>

  %2 = "tfl.reshape"(%arg0, %0) : (tensor<4x384x32x!quant.uniform<i8:f32, 0.2:-3>>, tensor<4xi32>) -> tensor<1x4x384x32x!quant.uniform<i8:f32, 0.2:-3>>
  // CHECK-NOT: tac.device tac.inference_type
  %3 = "tfl.quantize"(%2) {qtype = tensor<1x4x384x32x!quant.uniform<i8:f32, 0.19:1>>} : (tensor<1x4x384x32x!quant.uniform<i8:f32, 0.2:-3>>) -> tensor<1x4x384x32x!quant.uniform<i8:f32, 0.19:1>>

  // have side effects e.g. reduced flatbuffer size. Only certain type
  // conversions are supported.
  bool reduce_type_precision = false;
  // Whether to consider this model a quantized model with quantize/dequantize
  // ops and to convert kernels to quantized kernels wherever appropriate.
  quant::QDQConversionMode qdq_conversion_mode =
      quant::QDQConversionMode::kQDQNone;

  const std::unordered_set<std::string> tags_;

  const absl::flat_hash_map<std::string, tensorflow::SignatureDef>
      signature_def_map_;

  // Signature keys to identify the functions to load & quantize.
  const std::vector<std::string> signature_keys_;
};

// Runs passes to prepare the calibration model.
absl::Status RunCalibrationPasses(mlir::ModuleOp module_op, MLIRContext& ctx,

// The pass to decompose unregistered TF ops with the TFR compose function.
//
namespace mlir {
namespace TFR {

namespace {

// Quantize the float value based on given scale and zero point attributes.
IntegerAttr Quantize(float value, Attribute scale_attr, Attribute zp_attr,
                     OpBuilder builder) {
  double scale = mlir::cast<FloatAttr>(scale_attr).getValueAsDouble();

// whose storage type is 32-bit integer and expressed type is f32.
bool IsI32F32UniformQuantizedPerAxisType(Type type);

// Determines whether the storage type of a quantized type is supported by
// `tfl.quantize` or `tfl.dequantize` ops. ui8, i8 and i16 are supported.
bool IsSupportedByTfliteQuantizeOrDequantizeOps(IntegerType storage_type);

// Returns true if a type is quantized tensor type.
bool IsQuantizedTensorType(Type type);

  pm.addPass(TFL::CreatePostQuantizeRemoveQDQPass());
  if (failed(pm.run(module.get()))) {
    const std::string err(statusHandler.ConsumeStatus().message());
    LOG(ERROR) << "Failed to quantize: " << err;
    return kTfLiteError;
  }

  // Export the results.
  tflite::FlatbufferExportOptions options;
  options.toco_flags.set_force_select_tf_ops(false);

// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op -quant-prepare-quantize-drq='enable-per-channel-quantization=true' | FileCheck %s

module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>

  if (dot_general_op == nullptr) return std::nullopt;
  const int64_t filter_rank =
      mlir::dyn_cast<ShapedType>(dot_general_op.getOperand(1).getType())
          .getRank();

  // To quantize rhs per-channel, we currently only consider the case where
  // `stablehlo.dot_general` is legalizable to `tfl.fully_connected`.
  const bool is_per_axis_quantizable =
      IsDotGeneralFullyConnected(dot_general_op).value();

// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op -quant-prepare-quantize-drq | FileCheck %s

module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>