// CHECK-DAG: %[[w_quant:.*]] = "tf.Const"() <{value = dense<{{\[\[\[\[}}40, 20]
// CHECK-DAG: {{\[\[\[}}-87, -42]

// CHECK: %[[quantize:.*]] = "tf.PartitionedCall"(%arg0, %[[in_scale]], %[[in_zp]])
// CHECK-SAME: f = @quantize_i8
// CHECK: %[[conv_quant:.*]] = "tf.PartitionedCall"(%[[quantize]], %[[w_quant]], %[[b_quant]],
// CHECK-SAME: %[[in_scale]], %[[in_zp]], %[[w_scale]], %[[w_zp]],
// CHECK-SAME: %[[b_scale]], %[[w_zp]], %[[out_scale]], %[[out_zp]])

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

// CHECK-LABEL: func @conv_with_single_layer

// CHECK: %[[quantize:.*]] = "tf.PartitionedCall"(%arg0
// CHECK-SAME: f = @quantize_i8
// CHECK: %[[conv_quant:.*]] = "tf.PartitionedCall"(%[[quantize]]
// CHECK-SAME: f = @quantized_conv2d_with_bias_and_relu6_float_output_fn_0

    func.return %dot_out : tensor<*x!tf_type.qint32>
  }

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

    %act_max =  "tf.Const"() {value = dense<6.0> : tensor<f32>} : () -> tensor<f32>
    %i8_act_max_0 = "tf.PartitionedCall"(%act_max, %out_scale, %out_zp) {
        config = "", config_proto = "", executor_type = "", f=@quantize_i8
      } : (tensor<f32>, tensor<*xf32>, tensor<*xi32>) -> tensor<*xi8>
    %i8_act_max_1 = "tf.Cast"(%i8_act_max_0) {Truncate = false} : (tensor<*xi8>) -> tensor<*xf32>

// CHECK: %[[quantize_1:.*]] = "tf.PartitionedCall"(%arg0, %[[q_w]], %[[w_scale]], %[[w_zp]]) <{config = "", config_proto = "", executor_type = "", f = @quantized_conv2d_fn_1}> : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<*xf32>

    int bit_width = quant_specs_.GetQuantizationTypeWidth();

    Operation* quantize_op = quant_op.first;
    int quantize_operand_num = quant_op.second;

    auto affine_user = dyn_cast<AffineQuantizedOpInterface>(quantize_op);

    bool op_with_per_axis_support = false;

    if (!llvm::dyn_cast_or_null<CustomOp>(quantize_op)) {
      bool op_with_narrow_range =
          affine_user &&

    auto arg = bb.getArgument(0);

    auto remove_quantize_op = [&](QuantizeOp quantize_op) {
      auto quantize_output = quantize_op.getOutput();
      auto quantize_type = quantize_output.getType();
      input_types.push_back(quantize_type);
      auto new_arg = bb.addArgument(quantize_type, loc);
      quantize_output.replaceAllUsesWith(new_arg);
      quantize_op.erase();
      arg.dropAllUses();
      bb.eraseArgument(0);
    };

  }

  // Collects all candidate ops for quantization, which is the operand of
  // `quantize_op`. If successful, this always returns one element which is the
  // operand of `quantize_op`.
  FailureOr<SmallVector<Operation*>> CollectCandidateOps(
      QuantizeOpT quantize_op) const {
    Value operand = quantize_op->getOperand(0);
    if (QuantizedType::getQuantizedElementType(operand.getType())) {

  bool insertQDQ(PatternRewriter& rewriter, arith::ConstantOp op,
                 QuantizedType quant_type, QuantizationUnit quant_op) const {
    if (!quant_type) return false;

    Operation* quantize_op = quant_op.first;
    int quantize_operand_num = quant_op.second;

    Type expressed_type = op.getResult().getType();
    Type cast_type = quant_type.castFromExpressedType(expressed_type);

  auto test_func =
      module_op->lookupSymbol<func::FuncOp>("same_scale_after_composite");
  ASSERT_THAT(test_func, NotNull());

  auto quantize_op = FindOperationOfType<quantfork::QuantizeCastOp>(test_func);
  EXPECT_FALSE(IsOpQuantizableStableHlo(quantize_op));

  auto dequantize_op =
      FindOperationOfType<quantfork::DequantizeCastOp>(test_func);
  EXPECT_FALSE(IsOpQuantizableStableHlo(dequantize_op));
}