%7 = "tfl.quantize"(%6) {qtype = tensor<1x112x112x32x!quant.uniform<u8:f32, 1.0>>} : (tensor<1x112x112x32xf32>) -> tensor<1x112x112x32x!quant.uniform<u8:f32, 1.0>>
    func.return %7 : tensor<1x112x112x32x!quant.uniform<u8:f32, 1.0>>

// CHECK: %[[conv:.*]] = "tfl.conv_2d"(%arg0, %arg1, %arg2)
// CHECK-SAME: -> tensor<1x112x112x32x!quant.uniform<u8:f32, 0.0078431372549019607:128>>

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

// -----

func.func @testAvgPool(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
  // CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"

  func.return %3 : tensor<1x112x112x32xf32>

  func.return %3 : tensor<1x112x112x32xf32>

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

func.func @fully_connected(%arg0: tensor<1x37xf32>, %arg1: tensor<40x37xf32>, %arg2: tensor<40xf32>) -> tensor<1x40xf32> {
  // CHECK: _arithmetic_count = 3000 : i64

         exponential_avg_factor = 1.0 : f32,
         is_training = false
       }
        : (tensor<1x112x112x64xf32>, tensor<64xf32>, tensor<64xf32>, tensor<64xf32>, tensor<64xf32>)
       -> (tensor<1x112x112x64xf32>, tensor<64xf32>, tensor<64xf32>, tensor<64xf32>, tensor<64xf32>, tensor<64xf32>)

  func.return %2#0 : tensor<1x112x112x64xf32>
}

// CHECK-LABEL: func @fold_into_pad_with_extra_uses

// CHECK-SAME:          %[[INPUT:.*]]: tensor<1x112x112x64xf32>,
// CHECK-SAME:          %[[PAD_VAL:.*]]: tensor<f32>) -> tensor<1x64x114x114xf32> {
// CHECK:           %[[PAD:.*]] = stablehlo.pad %[[INPUT]], %[[PAD_VAL]],
// CHECK:               low = [0, 1, 1, 0], high = [0, 1, 1, 0], interior = [0, 0, 0, 0]
// CHECK:               : (tensor<1x112x112x64xf32>, tensor<f32>) -> tensor<1x114x114x64xf32>

// RUN: tf-opt %s -tf-optimize | FileCheck %s

// CHECK-LABEL: @fuseMulIntoConv2d
func.func @fuseMulIntoConv2d(%arg0: tensor<1x112x112x3xf32>) -> tensor<1x28x23x2xf32> {
  %cst0 = arith.constant dense<[[[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], [[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]], [[13.0, 14.0], [15.0, 16.0], [17.0, 18.0]]]]> : tensor<1x3x3x2xf32>
  %cst2 = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>

     %space_to_depth = "tf.SpaceToDepth"(%input) {block_size = 2, ...}:
        (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32>
     %device_launch = "tf_device.launch_func"(%space_to_depth,...) {func = @_func,...)
     return ...
   }
   func @_func(%input: tensor<2x112x112x12xf32>,
              %filter: tensor<7x7x3x64xf32>) {
     %filter_transform = "tf.Pad/tf.Transpose/tf.Reshape"(%filter):