#include "tensorflow/compiler/mlir/lite/utils/utils.h"

namespace mlir {
namespace TFL {
namespace tac {

bool NotTFLQuantDequantizeOp(Operation* op) {
  if (!op) return false;
  if (llvm::isa<TFL::QuantizeOp, TFL::DequantizeOp>(op)) return false;
  return true;
}

bool IsTerminatorOp(Operation* op) {
  if (!op) return false;
  return op->hasTrait<OpTrait::IsTerminator>();
}

// Try to guess the inference type of the op.

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

module {
  func.func @same_scale_test(%arg0: tensor<*xf32>) -> tensor<*xf32> {
    %cst = arith.constant dense<[-1, 144]> : tensor<2xi32>
    %cst_1 = arith.constant dense<1.0> : tensor<144x10xf32>
    %cst_2 = arith.constant dense<0.1> : tensor<10xf32>
    %0 = "quantfork.qcast"(%arg0) : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 0.05:-10>>

    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,

        ),
        tags={tag_constants.SERVING},
        signature_keys=['serving_default'],
    )

    model = quantize_model.quantize(
        temp_path,
        quantization_options=quantization_options,
        representative_dataset=data_gen(),
    )
    return model

  @test_util.run_in_graph_and_eager_modes

)

filegroup(
    name = "quantize_header",
    srcs = ["quantize.h"],
    visibility = ["//visibility:public"],
)

cc_library(
    name = "tfcompile_lib",
    srcs = [
        "codegen.cc",
        "compile.cc",
        "flags.cc",
    ],
    hdrs = [
        "codegen.h",
        "compile.h",
        "flags.h",
        "quantize.h",
    ],
    compatible_with = [],

  // It is known that `op` is `ModuleOp` when `pass` is
  // `QuantizeCompositeFunctionPass`, but the check is still performed to be
  // defensive.
  return pass != nullptr &&
         pass->getArgument() == "stablehlo-quantize-composite-functions" &&
         isa_and_nonnull<ModuleOp>(op);
}

// Report is saved only when:
// * After running `QuantizeCompositeFunctionPass`.
// * The pass is run on `ModuleOp`.

// RUN: tf-quant-opt %s -split-input-file -quant-prepare-quantize='post-training-quantize=true' | FileCheck %s

// -----

module {
  func.func @same_scale_ptq_test(%arg0: tensor<*xf32>) -> tensor<*xf32> {
    %cst = arith.constant dense<[-1, 144]> : tensor<2xi32>
    %cst_1 = arith.constant dense<1.0> : tensor<144x10xf32>
    %cst_2 = arith.constant dense<0.1> : tensor<10xf32>
    %0 = "quantfork.stats"(%arg0) {

// RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions='quantization-method=weight_only target-opset=XLA' -quant-quantize-composite-functions='quantization-method=weight_only target-opset=XLA enable-per-channel-quantization=true' -symbol-dce...

      if (QuantizedType::getQuantizedElementType(operand.getType())) {
        auto newTy = QuantizedType::castToExpressedType(operand.getType());
        newOperands.push_back(
            rewriter.create<TFL::DequantizeOp>(loc, newTy, operand));
        continue;
      }

      newOperands.push_back(operand);
    }

    SmallVector<Type> newResultTys;
    for (auto result : op->getResults()) {

  converters.reserve(dim_size);
  for (int i = 0, e = dim_size; i != e; ++i) {
    converters.push_back(getPerChunkConverter(i));
  }

  // Scan the elements of the dense elements attributes and quantize them by
  // using the right quantization parameters.
  int64_t flatten_index = 0;
  auto shape = type.getShape();
  int64_t chunk_size =
      std::accumulate(std::next(shape.begin(), quantization_dim_ + 1),