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>

  %dconv_s = "quantfork.stats"(%dconv) {layerStats = dense<[0.000000e+00, 1.000000e+01]> : tensor<2xf32>} : (tensor<1x112x112x64xf32>) -> tensor<1x112x112x64xf32>
  %bmm = "tfl.batch_matmul"(%conv_s, %dconv_s) {adj_x = false, adj_y = true} : (tensor<1x112x112x64xf32>, tensor<1x112x112x64xf32>) -> tensor<1x112x112x112xf32>

  func.return %conv : tensor<1x112x112x64xf32>

// CHECK: %[[b:.*]] = arith.constant dense<-1.23697901> : tensor<64xf32>
// CHECK: %[[w:.*]] = "tfl.pseudo_qconst"() <{qtype = tensor<64x3x3x3x!quant.uniform<i8<-127:127>:f32:0, {
// CHECK: %[[conv:.*]] = "tfl.conv_2d"(%arg0, %[[w]], %[[b]]) <{

  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

func.func @QuantizeFullyConnected(%arg0: tensor<1x224x224x3xf32>) -> tensor<1x112x112x4xf32> {
  %w = arith.constant dense<127.0> : tensor<4x12xf32>
  %b = arith.constant dense<0.0> : tensor<4xf32>
  %fc = "tfl.fully_connected"(%arg0, %w, %b) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x224x224x3xf32>, tensor<4x12xf32>, tensor<4xf32>) -> tensor<1x112x112x4xf32>
  func.return %fc : tensor<1x112x112x4xf32>

  return %3 : tensor<1x112x112x2xf32>
  // CHECK: %cst = arith.constant dense<0.000000e+00> : tensor<2xf32>
  // CHECK: %cst_0 = arith.constant dense<6.000000e+00> : tensor<1x3x3x2xf32>

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

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

  %4 = "tfl.dequantize"(%3) : (tensor<1x112x112x32x!quant.uniform<u8:f32, 1.0>>) -> tensor<1x112x112x32xf32>

  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