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

  %bmm = "tfl.batch_matmul"(%conv, %dconv) {adj_x = false, adj_y = true} : (tensor<1x112x112x64xf32>, tensor<1x112x112x64xf32>) -> tensor<1x112x112x112xf32>
  func.return %bmm, %emb : tensor<1x112x112x112xf32>, tensor<3x3x3x3xf32>

// CHECK-DAG: %[[b:.*]] = arith.constant dense<-1.23697901> : tensor<64xf32>

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


// -----

// CHECK-LABEL: testRandomUniform
func.func @testRandomUniform(%arg0: tensor<3xi32>) -> tensor<?x?x?xf32> {
  // CHECK: "tfl.random_uniform"(%arg0)

  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>

         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>

  %1 = "tfl.mul"(%0, %cst2) {fused_activation_function = "NONE"} : (tensor<1x112x112x2xf32>, tensor<1x112x112x2xf32>) -> tensor<1x112x112x2xf32>

  func.return %1 : tensor<1x112x112x2xf32>

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

    // CHECK-DAG: %[[SPACETODEPTH0:.*]] = "tf.SpaceToDepth"([[INPUT:.*]]) <{block_size = 2 : i64, data_format = "NHWC"}> : (tensor<2x224x224x3xf32>) -> tensor<2x112x112x12xf32>
    %2 = "tf.AddV2"(%arg2, %arg3) {device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    %3 = "tf.ReadVariableOp"(%arg6) : (tensor<!tf_type.resource<tensor<7x7x3x64xf32>>>) -> tensor<7x7x3x64xf32>