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

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

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

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

          &q_builder, input_model, quantized_type, use_updated_hybrid_scheme,
          ::tflite::optimize::QuantizerType::OLD_QUANTIZER) != kTfLiteOk) {
    return absl::InvalidArgumentError(
        "Quantize weights transformation failed.");
  }
  const uint8_t* q_buffer = q_builder.GetBufferPointer();
  *result =
      std::string(reinterpret_cast<const char*>(q_buffer), q_builder.GetSize());

// CHECK:           %[[VAL_20:.*]] = "tfl.quantize"(%[[VAL_19]]) <{qtype = tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>>}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x384x128xf32>) -> tensor<1x384x128x!quant.uniform<i8:f32, 3.000000e-01:-3>>

// CHECK-NOT: "tfl.batch_matmul"
func.func @Batchmatmul2FullyconnectedQDQ(%arg0: tensor<4x128x2xf32>, %arg1: tensor<2x1xf32>) -> (tensor<4x128x1xf32>) {
  %0 = arith.constant dense<[[1.0], [2.0]]> : tensor<2x1xf32>
  %1 = "tfl.quantize"(%0) {qtype = tensor<2x1x!quant.uniform<i8:f32, 0.024986599940879671:92>>} : (tensor<2x1xf32>) -> tensor<2x1x!quant.uniform<i8:f32, 0.024986599940879671:92>>

#include "absl/base/call_once.h"
#include "llvm-c/Target.h"
#include "llvm/Support/ManagedStatic.h"
#include "tensorflow/compiler/aot/codegen.h"
#include "tensorflow/compiler/aot/flags.h"
#include "tensorflow/compiler/aot/quantize.h"
#include "tensorflow/compiler/tf2xla/tf2xla.h"
#include "tensorflow/compiler/tf2xla/tf2xla_util.h"
#include "xla/client/client_library.h"
#include "xla/client/compile_only_client.h"
#include "xla/client/xla_computation.h"

// RUN: tf-quant-opt %s -split-input-file -quant-lift-quantizable-spots-as-functions -quant-quantize='target-opset=XLA' -verify-each=false | FileCheck %s

func.func private @conv(%input: tensor<1x3x4x3xf32> {tf._user_specified_name = "input_tensor"}) -> tensor<*xf32> attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<1x3x4x3>]} {
  %weight = arith.constant dense_resource<__elided__> : tensor<2x3x3x2xf32>

    readonly_model_ = input_model_->GetModel();
    model_ = UnPackFlatBufferModel(*readonly_model_);
  }
};

TEST_F(QuantizeLSTM2Test, VerifyLSTM) {
  // Quantize model.
  auto status = QuantizeModelAllOperators(
      &model_, TensorType_FLOAT32, TensorType_FLOAT32,
      /*allow_float=*/false, TensorType_INT8, output_buffer_);
  ASSERT_THAT(status, Eq(kTfLiteOk));