// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-prepare-quantize=enable-per-channel-quantized-weight=true -verify-diagnostics | FileCheck %s

// -----

module {
  // CHECK-LABEL: conv_with_bias_and_relu
  func.func private @conv_with_bias_and_relu(%arg0: tensor<1x3x2x3xf32>) -> tensor<1x2x2x2xf32> {
    %cst = "tf.Const"() {device = "", value = dense<[7.11401462, 7.05456924]> : tensor<2xf32>} : () -> tensor<2xf32>

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

// RUN: stablehlo-quant-opt %s -split-input-file -verify-diagnostics \
// RUN:     -stablehlo-quantize-composite-functions | FileCheck --check-prefix=CHECK %s

// Test that per-tensor weight-only quantized dot_general op is produced when
// empty `weight_only_ptq` is provided.

module attributes {tf_saved_model.semantics} {
  func.func private @quantize_dot_general_per_tensor(%arg0: tensor<1x2xf32>) -> tensor<1x3xf32> attributes {tf._original_func_name = "main_0"} {

// RUN: stablehlo-quant-opt %s -split-input-file -stablehlo-prepare-quantize=enable-per-channel-quantized-weight=false -verify-diagnostics | FileCheck %s

// -----

// CHECK-LABEL: func @dot
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<?x3xf32>) -> tensor<?x2xf32>
func.func @dot(%arg0: tensor<?x3xf32>) -> tensor<?x2xf32> {
  // CHECK: %[[cst:.*]] = stablehlo.constant
  // CHECK: %[[q1:.*]] = "quantfork.qcast"(%[[cst]])

    // asymmetric range. For a state tensor, assigning correct quantization
    // parameters is sufficient, and for constants with asymmetric range it's
    // not correctly quantized by legacy quantizer so call the new Quantize.
    return Quantize(real_value, tensor_type);
  } else if (width == 16) {
    if (const auto uniform_type = dyn_cast<UniformQuantizedType>(q_type)) {
      const auto quantized_values =

      quantization.}],
      "bool", "GetDynamicRangeQuantKernelSupport",
      (ins), [{}], [{return false;}]>,
    InterfaceMethod<
      [{Returns whether the op requires asymmetric quantize input attribute
      setting.}],
      "bool", "RequireAsymmetricQuantizeInputsAttr",
      (ins), [{}], [{return false;}]>,
  ];
}

// Specify this trait if the op has a fixed output value range.

        returned_type = quant::ConvertSignedQuantizedToUnsigned(
            dequantize_input.getType(), dequantize_op.getLoc());
        // replace the dequantize op by a quantize op
        TypeAttr type_attr = TypeAttr::get(returned_type);
        auto quantize_op = builder.create<QuantizeOp>(
            dequantize_op.getLoc(), returned_type, dequantize_input, type_attr);

class UsedBy<string op> : Constraint<
  CPred<"llvm::isa<mlir::TFL::" # op # "Op>(*$0.getUsers().begin())">>;

// When the op is passing-through, the output types of the quantized ops need
// to be updated as well. Since the quantize op manages its own type by the
// "qtype" attribute, we should update the type shape in this attribute.
def ReorderTransposeDequantQuant :
      Pat<(TF_TransposeOp:$old_value

// RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions='quantization-method=drq target-opset=UNIFORM_QUANTIZED' -quant-quantize-composite-functions='quantization-method=drq target-opset=UNIFORM_QUANTIZED' -symbol-dce | FileCheck %s

module {
  // TODO(b/260020937): Support transpose_a, transpose_b for matmul.
  func.func @matmul(%arg0: tensor<2x12xf32>) -> (tensor<*xf32>) {

    // `stablehlo.convolution` assumes the following format:
    // [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
    // `stablehlo.dot_general` can take various formats. We only per-channel
    // quantize non-batch ops.
    // `stablehlo.dot_general` legalizable to `tfl.fully_connected` has a
    // filter rank of 2 with the last dimension as the channel dimension.
    const int64_t quantization_dimension =