%cst_3 = "tf.Const"() {value = dense<[[[[1.400000e+01]], [[-2.800000e+01]], [[4.200000e+01]]], [[[-5.600000e+01]], [[7.100000e+01]], [[-8.500000e+01]]], [[[9.900000e+01]], [[-1.130000e+02]], [[1.270000e+02]]]]> : tensor<3x3x1x1xf32>} : () -> tensor<3x3x1x1xf32>
  %cst_4 = "tf.Const"() {value = dense<-1.280000e+02> : tensor<f32>} : () -> tensor<f32>
  %cst_5 = "tf.Const"() {value = dense<0.00118110236> : tensor<1xf32>} : () -> tensor<1xf32>

  return %2 : tensor<1x1x1x4xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_2(%arg0: tensor<1x1x1x8xf32>) -> tensor<1x1x1x4xf32>

// CHECK-LABEL: avg_pool2d_3
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>

  %3 = "tf.ReadVariableOp"(%arg2) : (tensor<!tf_type.resource<tensor<5xf32>>>) -> tensor<5xf32>
  %4 = "tf.ReadVariableOp"(%arg1) : (tensor<!tf_type.resource<tensor<3x3x1x5xf32>>>) -> tensor<3x3x1x5xf32>
  %5 = "tf.ReadVariableOp"(%arg3) : (tensor<!tf_type.resource<tensor<3920x10xf32>>>) -> tensor<3920x10xf32>

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>

        %0 = stablehlo.constant dense<2.000000e+00> : tensor<3x3x4x4xf32>
        %1 = stablehlo.convolution(%arg0, %0) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x3x4xf32>, tensor<3x3x4x4xf32>) -> tensor<1x3x3x4xf32>

// CHECK-LABEL: strips_attributes
// CHECK-NOT: tf
func.func @strips_attributes(%arg0: tensor<32x28x28x1xf32> {tf._user_specified_name = "x"},
                             %arg1: tensor<3x3x1x5xf32> {tf._user_specified_name = "w1"},
                             %arg2: tensor<5xf32> {tf._user_specified_name = "b1"},
                             %arg3: tensor<3920x10xf32> {tf._user_specified_name = "w2"},

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>

}

func.func @depthwiseConv2D(tensor<256x32x32x3xf32>, tensor<3x3x3x4xf32>, tensor<256x3x32x32xf32>) -> (tensor<256x30x30x12xf32>, tensor<256x12x30x30xf32>, tensor<256x30x30x12xf32>, tensor<256x30x30x12xf32>) {
^bb0(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<3x3x3x4xf32>, %arg2: tensor<256x3x32x32xf32>) :
   // OK

  %0 = "quantfork.stats"(%arg0) {
    layerStats = dense<[0.0, 1.0]> : tensor<2xf32>
  } : (tensor<?x5x5x2xf32>) -> tensor<?x5x5x2xf32>
  %1 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<3x5x5x2xf32>} : () -> tensor<3x5x5x2xf32>
  %2 = "tfl.pseudo_const"() {value = dense<0.000000e+00> : tensor<3xf32>} : () -> tensor<3xf32>
  %3 = "tfl.conv_2d"(%0, %1, %2) {
    dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32,

  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%in, %mini, %maxi) {num_bits = 4, narrow_range = true} : (tensor<3x3x3x4xf32>, tensor<4xf32>, tensor<4xf32>) -> tensor<3x3x3x4xf32>
  %rst = "tf.Conv2D"(%arg, %fq) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "SAME", strides = [1, 4, 5, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x4xf32>) -> tensor<256x8x7x4xf32>
  func.return %rst : tensor<256x8x7x4xf32>