// CHECK-SAME: tf_quant.quantized_ops = ["Conv3D", "BiasAdd"]
// CHECK: func private @quantized_batch_matmul_with_bias_fn
// CHECK-SAME: tf_quant.quantized_ops = ["BatchMatMul", "BiasAdd"]
// CHECK: func private @quantize_i8
// CHECK: func private @dequantize_i8

// UQ-CHECK-NOT: func private @internal_conv2d_fn
// UQ-CHECK-NOT: func private @internal_requantize_qi8_fn
// UQ-CHECK-NOT: func private @internal_requantize_no_activation_fn

// CHECK-NEXT: %[[V3:.*]] = "tf.PartitionedCall"(%[[V2]], %[[CST_4]], %[[CST_5]]) <{config = "", config_proto = "", executor_type = "", f = @dequantize_i8}> : (tensor<*xi8>, tensor<f32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: return %[[V3]] : tensor<*xf32>

// CHECK: func private @quantize_i8(
// CHECK: func private @dequantize_i8(

    %4 = "tf.PartitionedCall"(%3, %cst_2, %cst_3) {config = "", config_proto = "", executor_type = "", f = @dequantize_i8} : (tensor<1x2240x1120x512xi8>, tensor<f32>, tensor<i32>) -> tensor<1x2240x1120x512xf32>
    return %4 : tensor<1x2240x1120x512xf32>
  }
  func.func private @quantize_i8(%arg0: tensor<1x2240x2240x3xf32>, %arg1: tensor<f32>, %arg2: tensor<i32>) -> tensor<1x2240x2240x3xi8> {

// TF-DAG: %[[conv1_dequantized_1:.*]] = "tf.PartitionedCall"(%[[conv1_quantized]], %[[out_scale]], %[[in_out_zp]]) <{config = "", config_proto = "", executor_type = "", f = @dequantize_i8}>
// TF-DAG: %[[identity:.*]] = "tf.Identity"(%[[conv1_dequantized_1]])

// 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<*xf32>, tensor<*xf32>) -> tensor<*xf32>

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

  // Used for legacy 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>

                 METHOD_STATIC_RANGE_WEIGHT_ONLY_INT8) {
    // Uniform quantized opset is not supported for weight-only as inputs for
    // weight quantization are floats. And only dequantize_i8 is used from the
    // quantized function library.
    function_library_map = {
        {OpSet::TF, kQuantizedFunctionLibraryInMLIR},
        {OpSet::XLA, kQuantizedFunctionLibraryInMLIR_XLA_WEIGHT_ONLY}};
  } else {

    %2 = "tf.PartitionedCall"(%1, %cst_3, %cst_2) {config = "", config_proto = "", executor_type = "", f = @dequantize_i8} : (tensor<1x3x2x2xi8>, tensor<f32>, tensor<i32>) -> tensor<1x3x2x2xf32>
    return %2 : tensor<1x3x2x2xf32>
  }
  func.func private @quantize_i8(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<f32>, %arg2: tensor<i32>) -> tensor<1x3x4x3xi8> {

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