// the correct body. Only the type signatures matter for testing.
  func.func private @uniform_quantize_0(%arg0: tensor<1x2xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xi8>) -> tensor<1x2xi32> {
    %0 = stablehlo.convert %arg0 : (tensor<1x2xf32>) -> tensor<1x2xi32>
    return %0 : tensor<1x2xi32>
  }
// CHECK: @uniform_quantize_0

  // %[[RESULT0:.*]] = mhlo.uniform_dequantize %[[WHILE_RESULT]]#0
  %3 = "tf.UniformDequantize"(%2#0, %scales, %zps2) {quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64} : (tensor<2x?x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<2x?xf32>

  // %[[RESULT1:.*]] = mhlo.uniform_dequantize %[[WHILE_RESULT]]#0

    %0 = stablehlo.constant() {value = dense<127> : tensor<4x3xi8>} : () -> tensor<4x3x!quant.uniform<i8<-127:127>:f32:1, {5.000000e-03,5.000000e-03,5.000000e-03}>>
    %1 = stablehlo.uniform_quantize %arg0 : (tensor<1x4xf32>) -> tensor<1x4x!quant.uniform<i8:f32, 6.000000e-03:-128>>
    // CHECK: call @quantized_dot_general_relu_fn
    // CHECK-SAME: -> tensor<1x3xf32>

      %1 = stablehlo.uniform_quantize %arg0 : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>
      %2 = call @quantized_dot_general_fn(%1, %0) : (tensor<1x2x!quant.uniform<i8:f32, 4.000000e+0>>, tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {1.000000e+0,2.000000e+0,3.000000e+0}>>) -> tensor<1x3x!quant.uniform<i8:f32, 5.000000e+0>>

    // CHECK: %[[Q3:.*]] = stablehlo.uniform_quantize %0 : (tensor<1x3x!quant.uniform<i32:f32:1, {3.000000e-05,3.000000e-05,3.000000e-05}>>) -> tensor<1x3x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
    // CHECK: return %[[Q3]]

    // CHECK: %[[DOT:.*]] = stablehlo.dot_general
    // CHECK: %[[RQ:.*]] = stablehlo.uniform_quantize %[[DOT]]
    // CHECK: return %[[RQ]]

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>>
      return %0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
    }
  )mlir";

}

// -----

// CHECK-LABEL: func @uniform_quantize
func.func @uniform_quantize(%arg0: tensor<1xf32>) -> tensor<1x!tf_type.qint8>
{
  %scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>
  %zps = "tf.Const"() { value = dense<3> : tensor<i32> } : () -> tensor<i32>

  // CHECK: %[[qint:.*]] = "tf.UniformQuantize"

  %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>
  %1 = "tf.UniformDequantize"(%0, %scale, %zp) {
    quantization_axis = -1 : i64,
    quantization_min_val = -128 : 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,