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

}

def QuantConvertSimulatedQuant
    : Pass<"quant-convert-simulated-quantization", "func::FuncOp"> {
  let summary = "Converts training-time simulated quantization ops to "
                "corresponding quantize/dequantize casts";
  let constructor = "mlir::quantfork::createConvertSimulatedQuantPass()";
}

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

        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)

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

    } : (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",

      // Rewire the inputs.
      op->setOperand(index, duplicated_input_op->getResult(0));
    } else if (auto dq = dyn_cast<DequantizeOp>(input_op);
               dq && matchPattern(dq.getInput(), m_Constant(&attr))) {
      // Constant -> Dequantize case.
      builder->setInsertionPoint(op);
      Operation* duplicated_input_op =
          builder->clone(*dq.getInput().getDefiningOp());