// CHECK-NEXT:    return [[RES]], [[MATRES]] : tensor<1x1xf32>, tensor<1x1xf32>
//
module {
  func.func @serving_default(%arg0: tensor<1x3xf32>) -> (tensor<1x1xf32>, tensor<1x1xf32>) {
    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> tensor<3x1xf32>

  %outputs_2 = "tf.add"(%outputs_0, %outputs_1): (tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<1x1xf32>
  return %outputs_2 : tensor<1x1xf32>
}

// CHECK-LABEL: @_func
func.func private @_func(%arg0: tensor<1x3xf32>, %arg1: tensor<3x1xf32>) -> (tensor<1x1xf32>) {

  func.func @matmul(%arg0: tensor<2x12xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<12x2xf32>} : () -> tensor<12x2xf32>
    %1 = "tf.PartitionedCall"(%arg0, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn_1} : (tensor<2x12xf32>, tensor<12x2xf32>) -> tensor<*xf32>
    func.return %1: tensor<*xf32>
  }

  }

  // CHECK-LABEL: pack
  func.func @pack(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> {
    %0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
    func.return %0 : tensor<2x1xf32>
    // CHECK: %[[VAL_0:.*]] = arith.constant dense<[2, 1]> : tensor<2xi32>

  %min = arith.constant dense<0.0> : tensor<16xf32>
  %max = arith.constant dense<255.0> : tensor<16xf32>
  %mini = "tf.Identity"(%min) : (tensor<16xf32>) -> tensor<16xf32>
  %maxi = "tf.Identity"(%max) : (tensor<16xf32>) -> tensor<16xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%in, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<3x3x3x16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<3x3x3x16xf32>

    %4:4 = "tf.CustomAggregator"(%3) {calibration_method = 1 : i32, id = "1", max_percentile = 0.000000e+00 : f32, min_percentile = 0.000000e+00 : f32, num_bins = 0 : i32} : (tensor<3x11xf32>) -> (tensor<3x11xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
    %5 = stablehlo.add %4#0, %1 : tensor<3x11xf32>
    %6 = stablehlo.multiply %5, %1 : tensor<3x11xf32>

  %min = arith.constant dense<0.0> : tensor<16xf32>
  %max = arith.constant dense<15.0> : tensor<16xf32>
  %mini = "tf.Identity"(%min) : (tensor<16xf32>) -> tensor<16xf32>
  %maxi = "tf.Identity"(%max) : (tensor<16xf32>) -> tensor<16xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%in, %mini, %maxi) {num_bits = 3, narrow_range = false} : (tensor<3x3x3x16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<3x3x3x16xf32>

  %w = arith.constant dense<127.0> : tensor<4x12xf32>
  %b = arith.constant dense<0.0> : tensor<4xf32>
  %fc = "tfl.fully_connected"(%arg0, %w, %b) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x224x224x3xf32>, tensor<4x12xf32>, tensor<4xf32>) -> tensor<1x112x112x4xf32>
  func.return %fc : tensor<1x112x112x4xf32>

    %2 = "tf.Reshape"(%1, %cst_0) : (tensor<1x10xf32>, tensor<2xi32>) -> tensor<?x10xf32> loc(fused["Reshape:", "model/reshape"])
    %3 = "tf.AddV2"(%2, %cst) : (tensor<?x10xf32>, tensor<10xf32>) -> tensor<?x10xf32> loc(fused["AddV2:", "model/add"])
    func.return %3 : tensor<?x10xf32>
  }
}

// CHECK-LABEL: func @matmul_with_reshape_disable_quantization
// CHECK: "tf.PartitionedCall"

    %5 = stablehlo.constant dense<2.000000e+02> : tensor<1x1xf32>  // Output inverse scale.
    %6 = stablehlo.constant dense<2> : tensor<1x1xi8>  // Output zero point.
    %7 = call @uniform_quantize_0(%arg0, %0, %1) : (tensor<1x2xf32>, tensor<1x1xf32>, tensor<1x1xi8>) -> tensor<1x2xi8>
    %8 = stablehlo.convert %7 : (tensor<1x2xi8>) -> tensor<1x2xf32>