%zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>

  // CHECK: %[[qint:.*]] = "tf.UniformQuantize"
  // CHECK: %[[int:.*]] = "tf.Cast"(%[[qint]]) <{Truncate = false}> : (tensor<1x!tf_type.qint8>) -> tensor<1xi8>
  %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
    quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64

    // CHECK-DAG: %[[CST_1:.*]] = "tf.Const"() <{value = dense<3> : tensor<i32>}> : () -> tensor<i32>
    // CHECK-NEXT: %[[UQ:.*]] = "tf.UniformQuantize"(%arg0, %cst, %cst_0) <{quantization_axis = -1 : i64, quantization_max_val = 127 : i64, quantization_min_val = -128 : i64}> : (tensor<1xf32>, tensor<f32>, tensor<i32>) -> tensor<1x!tf_type.qint8>

  func.return %0 : tensor<1x3x3x2xf32>
}

//===----------------------------------------------------------------------===//
// tf.UniformQuantize and tf.UniformDequantize legalization
//===----------------------------------------------------------------------===//

// -----

// CHECK-LABEL: func @uniform_quantize_and_dequantize

      %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
        quantization_axis = -1 : i64, quantization_min_val = -2147483648 : i64, quantization_max_val = 2147483647 : i64
      } : (tensor<2x2xf32>, tensor<f32>, tensor<i32>) -> tensor<2x2x!tf_type.qint32>
      %qmax = "tf.UniformQuantize"(%max, %scales, %zps) {

      %scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
      %zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>

      %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
        quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64
      } : (tensor<?xf32>, tensor<f32>, tensor<i32>) -> tensor<?x!tf_type.qint8>

  %zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>

  // expected-error@+1 {{'tf.UniformQuantize' op is illegal as it is a UQ op or contains uq/qint types}}
  %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
    quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64

func.func @main(%arg0: tensor<10xf32>) -> tensor<10xf32> {
  %scale = "tf.Const"() { value = dense<0.347> : tensor<f32> } : () -> tensor<f32>
  %zp = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>
  %0 = "tf.UniformQuantize"(%arg0, %scale, %zp) {
    quantization_axis = -1 : i64,
    quantization_min_val = -128 : i64,
    quantization_max_val = 127 : i64
  } : (tensor<10xf32>, tensor<f32>, tensor<i32>) -> tensor<10x!tf_type.qint8>

  // CHECK-DAG: %[[LHS2:.*]] = mhlo.bitcast_convert %[[LHS1]] : (tensor<3x2xi32>) -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
  %0 = "tf.UniformQuantize"(%input, %input_scales, %input_zps) {
    quantization_axis = -1 : i64, quantization_min_val = -2147483648 : i64, quantization_max_val = 2147483647 : i64

  // Quantize initial input at the start of the graph. Output is qint8.
  func.func @quantize_i8(%input : tensor<*xf32>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>) -> tensor<*x!tf_type.qint8> {
    %quantize = "tf.UniformQuantize"(%input, %input_scale, %input_zp) {
      Tin = "tfdtype$DT_FLOAT",
      Tout = "tfdtype$DT_QINT8",
      quantization_axis = -1,
      quantization_min_val = -128,
      quantization_max_val = 127,

  for (auto [arg, arg_type, arg_loc] :
       llvm::zip_equal(entry.getArguments(), arg_types, arg_locs)) {
    arg.setType(arg_type);
    arg.setLoc(arg_loc);
  }
}

// Creates a UniformQuantize op and sets it as return op.
// The requantize scale and zero point should be determined from the
// `entry_func_op`'s output, containing information on layerStats of the
// entire function.