// CHECK: %[[dequantize:.*]] = "tf.PartitionedCall"(%[[conv_quant]], %[[out_scale]], %[[out_zp]])
// CHECK-SAME: f = @dequantize_i8

// CHECK: %[[conv_float:.*]] = "tf.PartitionedCall"(%arg0, %[[w_float]], %[[b_float]])
// CHECK-SAME: f = @composite_conv2d_with_bias_and_relu6_fn_1

// CHECK: return %[[dequantize]], %[[conv_float]]

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

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

  // For weight-only
  func.func @dequantize_i8(%input : tensor<*xi8>, %scale : tensor<*xf32>, %zp : tensor<*xi32>) -> tensor<*xf32> {
    // Use identity op to avoid the weight being constant-folded.
    %identity = "tf.Identity"(%input) : (tensor<*xi8>) -> tensor<*xi8>

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

// CHECK: %[[dequantize:.*]] = "tf.PartitionedCall"(%[[maxpool]]
// CHECK-SAME: f = @dequantize_i8
// CHECK: return %[[dequantize]]

// CHECK: -------- Quantization Summary --------
// CHECK: Number of quantized layers in the model
// CHECK: --------------------------------
// CHECK: Name    Count/Total

    %round = "tf.Round"(%clamp_min) : (tensor<*xf32>) -> tensor<*xf32>
    %i8 = "tf.Cast"(%round) : (tensor<*xf32>) -> tensor<*xi8>
    func.return %i8 : tensor<*xi8>
  }

  func.func @dequantize_i8(%input : tensor<*xi8>, %scale : tensor<*xf32>, %zp : tensor<*xi32>) -> tensor<*xf32> {
    // Use identity op to avoid the weight being constant-folded.
    %identity = "tf.Identity"(%input) : (tensor<*xi8>) -> tensor<*xi8>

      : OpRewritePattern<RootOpT>(context, /*benefit=*/300) {}

 private:
  // Collects all candidate ops for quantization, which are the
  // `dequantize_op`'s users.
  FailureOr<SmallVector<Operation*>> CollectCandidateOps(
      DequantizeOpT dequantize_op) const {
    auto users = dequantize_op->getResult(0).getUsers();
    return SmallVector<Operation*>(users.begin(), users.end());
  }

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

TEST_F(IsOpQuantizableStableHloTest,
       XlaCallModuleOpQuantizableWhenNotDenylisted) {

    Operation* returned_op = returned_value.getDefiningOp();
    if (returned_op && returned_op->hasOneUse() &&
        llvm::isa<DequantizeOp>(returned_op)) {
      auto dequantize_op = llvm::cast<DequantizeOp>(returned_op);
      Value dequantized_result = dequantize_op.getInput();
      output_types.push_back(dequantized_result.getType());
      terminator->setOperand(i, dequantized_result);
      returned_op->erase();
    } else {