// QDQ-NEXT:  return %[[split]]#0 : tensor<2xf32>
}

// CHECK-LABEL: FoldTranspose
func.func @FoldTranspose(%arg0: tensor<1x10x20x3xf32>) -> tensor<1x20x40x16xf32> {
  %cst = arith.constant dense<[1, 20, 40, 16]> : tensor<4xi32>
  %cst_0 = arith.constant dense<[2, 0, 1, 3]> : tensor<4xi32>

module {
  func.func @wrap_composite_func(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {
    %0 = "tf.PartitionedCall"(%arg0, %arg1) <{f = @composite_conv2d_with_relu6_fn}> {_tfl_quant_trait = "fully_quantizable"}
          : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    func.return %0 : tensor<*xf32>
  }

  func.func @no_composite_func(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> {

  %custom_2 = "tfl.custom"(%0, %w_1, %w_2, %b) {custom_code = "CustomTestOp2", custom_option = #tfl<const_bytes : "0x">} : (tensor<1x1x1x1xf32>, tensor<4096x1x1x1xf32>, tensor<128x1x1x1xf32>, tensor<2048x1x1x1xf32>) -> tensor<*xf32>

func.func @uint8_to_int8(%arg0: tensor<2x2xf32>) -> tensor<2x2xf32> {
  %1 = "tfl.quantize"(%arg0) {qtype = tensor<2x2x!quant.uniform<u8:f32, 1.0:128>>} : (tensor<2x2xf32>) -> tensor<2x2x!quant.uniform<u8:f32, 1.0:128>>
  %2 = "tfl.dequantize"(%1) : (tensor<2x2x!quant.uniform<u8:f32, 1.0:128>>) -> tensor<2x2xf32>
  func.return %2 : tensor<2x2xf32>

  } : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
  return %0 : tensor<1x3xf32>
}

// CHECK-LABEL: func.func @qdq_for_dot_general_weight_empty
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x2xf32>) -> tensor<1x3xf32>
// CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3xf32>}> : () -> tensor<2x3xf32>

    %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<1x3xf32>) -> (tensor<1x3xf32>, tensor<f32>, tensor<f32>, tensor<*xi64>)
    %5 = stablehlo.constant dense<1.000000e+03> : tensor<3x64xf32>

  func.return %0 : tensor<4xf32>
}

// CHECK:       func @squaredDifference(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
// CHECK:         %0 = "tf.Sub"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// CHECK:         %1 = "tf.Mul"(%0, %0) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
// CHECK:         return %1 : tensor<4xf32>
// CHECK:       }

// -----

    "tfl.yield"(%6, %5, %32, %17, %20, %0) : (tensor<i32>, tensor<i32>, tensor<*xf32>, tensor<4x2xf32>, tensor<4x2xf32>, tensor<4x4x3xf32>) -> ()
  }) {is_stateless = true} : (tensor<i32>, tensor<i32>, tensor<4x4x2xf32>, tensor<4x2xf32>, tensor<4x2xf32>, tensor<4x4x3xf32>) -> (tensor<i32>, tensor<i32>, tensor<*xf32>, tensor<4x2xf32>, tensor<4x2xf32>, tensor<*xf32>)

  %2 = mhlo.constant dense<0.000000e+00> : tensor<f32>
  %3 = "mhlo.pad"(%1, %2) <{edge_padding_high = dense<1> : tensor<3xi64>, edge_padding_low = dense<0> : tensor<3xi64>, interior_padding = dense<0> : tensor<3xi64>}> : (tensor<11x11x11xf32>, tensor<f32>) -> tensor<12x12x12xf32>
  return %3 : tensor<12x12x12xf32>

    R"mlir(
  func.func @main(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
    %if = "stablehlo.if"(%arg0) ({
      %0 = stablehlo.add %arg1, %arg1 : tensor<f32>
      stablehlo.return %0 : tensor<f32>
    }, {
      %1 = stablehlo.add %arg2, %arg2 : tensor<f32>
      stablehlo.return %1 : tensor<f32>
    }) : (tensor<i1>) -> (tensor<f32>)