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>
// CHECK: return %[[CUSTOM_0]] : tensor<1x2x2x1xf32>
}

         padding = "EXPLICIT",
         strides = [5, 6, 7, 8]
       } : (tensor<1x32x32x3xf32>, tensor<1x1x3x8xf32>) -> tensor<1x7x7x8xf32>

  func.return %0 : tensor<1x7x7x8xf32>
}

// CHECK-LABEL: func @transposeConv2DWithDefaultAttr
func.func @transposeConv2DWithDefaultAttr(%input: tensor<1x32x32x3xf32>, %filter: tensor<1x1x3x8xf32>) -> tensor<?x?x?x?xf32>
{

  %conv2 = "tfl.conv_2d"(%0, %w, %b2) {
    dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU",
    padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32
  } : (tensor<1x5x5x2xf32>, tensor<3x1x1x2xf32>, tensor<3xf32>) -> tensor<1x5x5x3xf32>
  func.return %conv, %conv2 : tensor<1x5x5x3xf32>, tensor<1x5x5x3xf32>

// RUN: tf-quant-opt %s -split-input-file -quant-propagate-quantize-type | FileCheck %s

module {
  func.func @not_propagate_matmul(%arg0: tensor<1x2x2x2xf32>) -> tensor<*xf32> {
    %cst = "tf.Const"() {value = dense<127> : tensor<2x1024xi8>} : () -> tensor<2x1024xi8>
    %cst_0 = "tf.Const"() {value = dense<0.0157480314> : tensor<f32>} : () -> tensor<f32>
    %0 = "tf.Identity"(%cst) : (tensor<2x1024xi8>) -> tensor<2x1024xi8>

// CHECK-LABEL: QuantizeUnidirectionalLstmFullPerTensor
func.func @QuantizeUnidirectionalLstmFullPerTensor(%arg0: tensor<1x2x3xf32>) -> (tensor<1x2x3xf32>) {
  %input = "quantfork.stats"(%arg0) {layerStats = dense<[0.0, 1.0]> : tensor<2xf32>} : (tensor<1x2x3xf32>) -> tensor<1x2x3xf32>
  %1 = "tfl.pseudo_const"() {value = dense<[[0.1]]> : tensor<1x1xf32>} : () -> tensor<1x1xf32>

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

func.func private @conv(%input: tensor<1x3x4x3xf32> {tf._user_specified_name = "input_tensor"}) -> tensor<*xf32> attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<1x3x4x3>]} {
  %weight = arith.constant dense_resource<__elided__> : tensor<2x3x3x2xf32>
  %bias = arith.constant dense<[7.11401462, 7.05456924]> : tensor<2xf32>

  %1 = "tf.BiasAdd"(%0, %cst_0) {data_format = "NHWC"} : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32>
  %2 = "tf.Mul"(%0, %cst_1) : (tensor<1x3x2x2xf32>, tensor<2xf32>) -> tensor<1x3x2x2xf32>
  func.return %1, %2 : tensor<1x3x2x2xf32>, tensor<1x3x2x2xf32>
}
// CHECK: func @not_fuse_conv2d_with_bias_and_mul

}

// -----

// CHECK-LABEL: testReverseV2
func.func @testReverseV2(%arg0: tensor<1x2x3x4xf32>, %arg1 : tensor<2xi32>) -> tensor<1x2x3x4xf32> {
  // CHECK: "tfl.reverse_v2"(%arg0, %arg1)
  %0 = "tfl.reverse_v2"(%arg0, %arg1): (tensor<1x2x3x4xf32>, tensor<2xi32>) -> tensor<1x2x3x4xf32>
  func.return %0 : tensor<1x2x3x4xf32>
}

// -----

// test select
// CHECK-LABEL: testSelect

    %2 = "quantfork.stats"(%1) {layerStats = dense<[5.00000000e-6, 7.00000000e-1]> : tensor<2xf32>} : (tensor<2x2x3xf32>) -> tensor<2x2x3xf32>
    return %2 : tensor<2x2x3xf32>
  }
// CHECK: func.func private @quantize_dot_general_batch_per_tensor_quantized_fn(%[[ARG_0:.+]]: tensor<2x2x2xf32>) -> tensor<2x2x3xf32> attributes {tf._original_func_name = "main_0"}