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:               : (tensor<1x114x114x64xf32>) -> tensor<1x64x114x114xf32>
// CHECK:           return %[[TPOS]] : tensor<1x64x114x114xf32>

func.func @commute_transpose_pad(
      %arg0: tensor<1x112x112x64xf32>, %padding_val: tensor<f32>)
      -> tensor<1x64x114x114xf32> {
  %tspos = stablehlo.transpose %arg0, dims = [0, 3, 1, 2]
      : (tensor<1x112x112x64xf32>) -> tensor<1x64x112x112xf32>

  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

    %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  %3, %6 : tensor<?xf32>, tensor<?xf32>
}

// -----

// CHECK-LABEL: func @skip_nan_inf_constant
// CHECK-SAME: (%[[ARG_0:.*]]: tensor<?x112x112x64xf32>) -> tensor<?x56x56x64xf32>
func.func @skip_nan_inf_constant(%arg0: tensor<?x112x112x64xf32>) -> tensor<?x56x56x64xf32> {
  // CHECK-DAG: %[[cst0:.*]] = stablehlo.constant dense<0xFF800000> : tensor<f32

  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>

// 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>

               %filter: tensor<7x7x3x64xf32>) {
     %6 = "tf.Conv2D"(%input, %filter)  {strides = [1, 2, 2, 1]}:
                 (tensor<2x230x230x3xf32>, tensor<7x7x3x64xf32>) ->
      tensor<2x112x112x64xf32>
   }
}

// With this pass, the program will be transformed into:
module {
   func @while_body {
     %input = "tf.IteratorGetNext"(...) {device = "/CPU:0"}
               -> tensor<2x224x224x3xf32>