// CHECK: %3 = tfl.mul %2, %cst_1 {fused_activation_function = "NONE"} : tensor<1x2x2x1xf32>
// CHECK: %cst_2 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>
// CHECK: %4 = "tfl.transpose"(%3, %cst_2) : (tensor<1x2x2x1xf32>, tensor<4xi32>) -> tensor<1x1x2x2xf32>
// CHECK: return %4 : tensor<1x1x2x2xf32>

func.func @avg_pool2d_6(%arg0: tensor<1x1x1x7xf32>) -> (tensor<1x1x1x2xf32>) {

func.func @depth_to_space(%arg0: tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32> {
  %0 = "tf.DepthToSpace"(%arg0) {block_size = 2: i64,  data_format = "NHWC"}: (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
  func.return %0 : tensor<1x2x2x1xf32>
// CHECK: %[[CUSTOM_0:.*]] = "tfl.custom"(%arg0) <{custom_code = "FlexDepthToSpace", custom_option = #tfl<const_bytes : "{{.*}}">}> : (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>

    %3 = "tf.Identity"(%2) {device = ""} : (tensor<1x3x1x1xf32>) -> tensor<1x3x1x1xf32>
    return %3 : tensor<1x3x1x1xf32>
  }

// CHECK-LABEL: func @multiple_quantizable_ops_in_graph

    : (tensor<?x3x2x1xf32>, tensor<2x1x1x1xf32>) -> tensor<?x2x2x1xf32>
  %1 = chlo.broadcast_add %0, %zp_offset : (
      tensor<?x2x2x1xf32>, tensor<?x2x2x1xf32>) -> tensor<?x2x2x1xf32>
  %2 = chlo.broadcast_add %1, %bias : (
      tensor<?x2x2x1xf32>, tensor<1xf32>) ->tensor<?x2x2x1xf32>
  return %2 : tensor<?x2x2x1xf32>
}

// -----

// CHECK-LABEL: func @dot_general_add_add

    %1 = stablehlo.constant dense<1.000000e+03> : tensor<1x1x1x1xf32>  // Input inverse scale.
    %2 = stablehlo.constant dense<-128> : tensor<1x1x1x1xi8>  // Input zero point.
    %3 = stablehlo.constant dense<1> : tensor<3x3x4x4xi8>  // Quantized filter tensor.
    %4 = stablehlo.constant dense<3.000000e+03> : tensor<1x1x1x4xf32>
    %5 = stablehlo.constant dense<4.000000e+03> : tensor<1x1x1x1xf32>  // Output inverse scale.

                                      interior_padding = #vhlo.tensor_v1<dense<0> : tensor<3xi64>>}> : (tensor<1x160x1xf32>, tensor<f32>) -> tensor<1x161x1xf32>
  return %0 : tensor<1x161x1xf32>
}

//CHECK:func.func private @pad(%arg0: tensor<1x160x1xf32>, %arg1: tensor<f32>) -> tensor<1x161x1xf32> {
//CHECK-NEXT: %0 = "vhlo.pad_v1"(%arg0, %arg1) <{edge_padding_high = #vhlo.tensor_v1<dense<0> : tensor<3xi64>>,

// MinElement-LABEL: QuantizeCustomOp
func.func @QuantizeCustomOp(%arg0: tensor<1x1x1x1xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>) attributes {tf.entry_function = {inputs = "input", outputs = "custom_op"}} {
  %0 = "quantfork.stats"(%arg0) {layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>} : (tensor<1x1x1x1xf32>) -> tensor<1x1x1x1xf32>
  %w_1 = arith.constant dense<127.0> : tensor<4096x1x1x1xf32>

// CHECK:           %[[VAL_6:.*]] = "tfl.reshape"(%[[VAL_1]], %[[VAL_2]]) : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
// CHECK:           %[[VAL_7:.*]] = "tfl.concatenation"(%[[VAL_5]], %[[VAL_6]]) <{axis = 3 : i32, fused_activation_function = "NONE"}> : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>

    %1 = "tfl.reshape"(%arg1, %cst) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>
    %2 = "tfl.concatenation"(%0, %1) {axis = 3 : i32, fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>

  %output = "tf.Cast"(%q_output) {Truncate = false} : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32>

  // CHECK-DAG: %[[OUTPUT_2:.*]] = "tf.Cast"(%[[OUTPUT_QINT]]) <{Truncate = false}> : (tensor<1x2x2x1x!tf_type.qint32>) -> tensor<1x2x2x1xi32>
  // CHECK-DAG: %[[OUTPUT_QINT_1:.*]] = "tf.Cast"(%[[OUTPUT_1]]) <{Truncate = false}> : (tensor<1x2x2x1xi32>) -> tensor<1x2x2x1x!tf_type.qint32>