// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-prepare-quantize=bit-width=4 -verify-diagnostics | FileCheck %s

// CHECK-LABEL: func @dot_int4
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<?x3xf32>) -> tensor<?x2xf32>
func.func @dot_int4(%arg0: tensor<?x3xf32>) -> tensor<?x2xf32> {
  // CHECK: %[[cst:.*]] = stablehlo.constant
  // CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[cst]])
  // CHECK-SAME: quant.uniform<i8:f32, 0.0040316890267764818:127>

  EXPECT_FALSE(IsOpFullyQuantized(*add_op_itr));
}

TEST_F(IsOpFullyQuantizedTest, FalseIfOpPartiallyQuantized) {
  constexpr absl::string_view kQuantizeOp = R"mlir(
    func.func @quantize(%arg0: tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>> {
      %0 = stablehlo.uniform_quantize %arg0 : (tensor<2xf32>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>

    // 1. Collect quantizable ops.
    QuantizationUnits quantizable_ops = GetQuantizableOps(op);
    if (quantizable_ops.empty()) {
      return failure();
    }

    // 2. Quantize collected ops.
    if (!QuantizeOps(rewriter, op, quantizable_ops)) {
      return failure();
    }

    // 3. Complete the Q-DQ pair for each inference type.
    if (!ConvertToFloat16Constant(rewriter, op)) {

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

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