include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"

// Quantize attribute $0 by using quantization parameter from %1.
def QuantizeByQuantizedType : NativeCodeCall<"quant::Quantize($0, $1.getValue())">;
def F32ElementsAttr : ElementsAttrBase<
  CPred<"$_self.cast<ElementsAttr>().getShapedType().getElementType().isF32()">, "float constant tensor">;

// Squash tfl.dequantize and tfl.quantize pairs.

// Stores information about how to quantize a user-specified custom operation.
// CustomOpInfo contains info of its corresponding CustomOp registered in the
// CustomOpMap. 'quantizable_input_indices' is used to determine which indices
// of the CustomOp are quantizable. 'is_weight_only' is used specify whether the
// custom op is quantized only for storage and dequantized at runtime.

                  const tensorflow::quantization::PyFunctionLibrary*
                      quantization_py_function_library = nullptr);

// Quantize the model with calibration data. Throw errors if `fully_quantize`
// is specified by the calibration data are not sufficient to quantize the
// model.
PyObject* MlirQuantizeModel(PyObject* data, bool disable_per_channel,
                            bool fully_quantize, int inference_type,

    enable_whole_model_verify,
    denylisted_ops,
    denylisted_nodes,
    enable_variable_quantization,
    disable_per_channel_for_dense_layers,
    debug_options_str,
):
  """Wraps experimental mlir quantize model."""
  return _pywrap_converter_api.ExperimentalMlirQuantizeModel(
      input_data_str,
      disable_per_channel,
      fully_quantize,
      inference_type,
      input_data_type,
      output_data_type,

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

  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,

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

      mlir::TFL::PassConfig(quant_specs), pm);

  if (failed(pm.run(module.get()))) {
    absl::string_view err = statusHandler.ConsumeStatus().message();
    LOG(ERROR) << "Failed to quantize: " << err;
    return kTfLiteError;
  }

  // Export the results to the builder
  std::string result;
  tflite::FlatbufferExportOptions options;
  options.toco_flags.set_force_select_tf_ops(false);