module attributes {tf.versions = {producer = 179 : i32}} {
  func.func @main(%arg0: tensor<f32>) -> tensor<f32> {
    %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<ui8:f32, 34.0:16>>
    %1 = "stablehlo.uniform_dequantize"(%0) : (tensor<!quant.uniform<ui8:f32, 34.0:16>>) -> tensor<f32>
    func.return %1 : tensor<f32>
  }
}

// CHECK-LABEL: HloModule main

      "uniform_quantize". The first argument is the tensor to be quantized,
      the second argument is the zero point constant (element type: int) and
      the third argument is the inverse scale constant (element type: float).
    * A tensor is dequantized using a `func::FuncOp` whose name contains
      "uniform_dequantize". The first argument is the tensor to be quantized,

        return success();
      }

      op.replaceAllUsesWith(q.getArg());
      return success();
    }
    return failure();
  }
};

// Replaces constant and uniform_quantize ops with single quantized constant op.
class QuantizeConstPattern
    : public OpRewritePattern<mlir::stablehlo::UniformQuantizeOp> {
 public:
  explicit QuantizeConstPattern(MLIRContext* context)

// CHECK-NO-UNPACK: %[[QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x1024xf32>) -> tensor<1x1024x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-NO-UNPACK: %[[DOT:.+]] = stablehlo.dot_general %[[QUANTIZE_0]], %[[CONST]]
// CHECK-NO-UNPACK: %[[QUANTIZE_1:.+]] = stablehlo.uniform_quantize %[[DOT]] : (tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>>) -> tensor<1x3x!quant.uniform<i8:f32, {{.*}}>>

  %0 = stablehlo.uniform_quantize %arg : (tensor<2x2x!quant.uniform<i8:f32, 2.000000e+0:16>>) -> tensor<2x2x!quant.uniform<i8:f32, 3.000000e+0:127>>
  return %0 : tensor<2x2x!quant.uniform<i8:f32, 3.000000e+0:127>>
}
// CHECK-LABEL: uniform_quantize_op_quantized_input
// CHECK: stablehlo.uniform_quantize
// CHECK-NOT: tfl.quantize

// -----

    // 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";

  %output_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
  %output_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>

  // CHECK-DAG: %[[LHS:.*]] = mhlo.uniform_quantize %arg0 : (tensor<3x2xf32>) -> tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>
  // CHECK-DAG: %[[LHS1:.*]] = mhlo.bitcast_convert %[[LHS]] : (tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>) -> tensor<3x2xi32>

}

// -----

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