// CHECK: %[[DEQUANTIZE:.*]] = mhlo.uniform_dequantize %[[CONVERT_2]] : (tensor<2x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<2xf32>
  // CHECK: return %[[DEQUANTIZE]] : tensor<2xf32>

  %0 = "tf.UniformQuantize"(%arg0, %scales, %zps) {
    quantization_axis = -1 : i64, quantization_min_val = -128 : i64, quantization_max_val = 127 : i64

      } : (tensor<i8>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xf32>

    %clamp_max = "tf.Maximum"(%dequantize, %clip_min) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    %clamp_min = "tf.Minimum"(%clamp_max, %clip_max) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>

    func.return %clamp_min : tensor<*xf32>
  }

  // Dequantizes and applies quantized Relu by clipping.

// CHECK-LABEL: fuseMulIntoPerTensorConv2dWithQDQs
func.func @fuseMulIntoPerTensorConv2dWithQDQs(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x8x7x3xf32> {
  %cst = arith.constant dense<1.5> : tensor<3xf32>
  %cst_0 = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
  %w = arith.constant dense<2.0> : tensor<3x3x3x3xf32>

        // correct float op should be the user of the last DequantizeOp.
        if (llvm::isa<QuantizeOpT>(user)) {
          user = *user->getResult(0).getUsers().begin();
        }
        if (auto dequantize = llvm::dyn_cast<DequantizeOpT>(user)) {
          // Replace all uses, except not quantizable ops that are being used in
          // the float backbone.

  %6 = "tfl.dequantize"(%5) : (tensor<1x401408x!quant.uniform<i8:f32, 3.906250e-03>>) -> tensor<1x401408xf32>
  func.return %6 : tensor<1x401408xf32>

// CHECK-LABEL: func @modified(%arg0: tensor<1x224x224x3xf32>) -> tensor<1x401408xf32>
// CHECK-NEXT: %[[shape:.*]] = arith.constant dense<[1, 401408]> : tensor<2xi32>

  // CHECK-DAG:  %[[CST_1:.*]] = "tfl.pseudo_const"() <{value = dense<0.000000e+00> : tensor<4x4xf16>}> : () -> tensor<4x4xf16>
  // CHECK-DAG:  %[[DQ_0:.*]] = "tfl.dequantize"(%[[CST_0]]) : (tensor<4x4xf16>) -> tensor<4x4xf32>
  // CHECK-DAG:  %[[DQ_1:.*]] = "tfl.dequantize"(%[[CST_1]]) : (tensor<4x4xf16>) -> tensor<4x4xf32>

  // TODO: b/323478683 - Make the attribute being part of op definition.
  quantize->setAttr(kVolatileOpAttrName, builder_.getUnitAttr());

  // `original_result` has a use to `quantize`, so this will replace that use
  // by the result of `dequantize`. Remember to reset that use afterwards
  value.replaceAllUsesWith(dequantize);
  quantize.getOperation()->replaceUsesOfWith(dequantize, value);
}

// `tfl.quantize` or `tfl.dequantize` ops. ui8, i8 and i16 are supported.
bool IsSupportedByTfliteQuantizeOrDequantizeOps(IntegerType storage_type) {
  if (storage_type.getWidth() == 8 ||
      (storage_type.isSigned() && storage_type.getWidth() == 16)) {
    return true;
  }
  LLVM_DEBUG(llvm::dbgs()
             << "Uniform quantize / dequantize op only supports ui8, i8 or "

        mlir::cast<ShapedType>(call_op.getResult(0).getType())
            .clone(rewriter.getF32Type()));
    rewriter.setInsertionPoint(call_op);
    rewriter.insert(new_call_op);

    // Remove the dequantize ops and replace uses by the new func.call op.
    SmallVector<Operation*> users_to_erase;
    for (auto user : users) {
      llvm::dyn_cast<mlir::stablehlo::UniformDequantizeOp>(user)

    } : (tensor<*xf32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*x!tf_type.qint32>
    func.return %quantize : tensor<*x!tf_type.qint32>
  }

  // Dequantize final graph output back to f32. Input is qint8.
  func.func @dequantize_i8(%input : tensor<*x!tf_type.qint8>, %input_scale : tensor<*xf32>, %input_zp : tensor<*xi32>) -> tensor<*xf32> {
    %dequantize = "tf.UniformDequantize"(%input, %input_scale, %input_zp) {
      Tin = "tfdtype$DT_QINT8",