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

// UQ-CHECK: func private @quantized_matmul_with_relu_fn
// UQ-CHECK: func private @quantized_matmul_with_relu6_fn
// UQ-CHECK: func private @quantize_i8
// UQ-CHECK: func private @quantize_i32

    %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> {

// CHECK-NO-UNPACK: %[[QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x1024xf32>) -> tensor<1x1024x!quant.uniform<i8:f32, {{.*}}>>
// CHECK-NO-UNPACK: %[[DOT:.+]] = stablehlo.dot_general %[[QUANTIZE_0]], %[[CONST]]

        quantize_op.setOperand(new_arg);
      } else {
        input_type.print(llvm::errs() << "Requested input type ");
        quantize_op.emitError(" Couldn't be modified to the requested type.");
        return failure();
      }
      new_input_types[i] = arg_type;
      arg.dropAllUses();
      if (quantize_op.use_empty()) {
        quantize_op.erase();
      }
    } else {

                 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 {

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