%1 = "tf.BiasAdd"(%0, %cst) {data_format = "NHWC", device = ""} : (tensor<1x2x2x3xf32>, tensor<3xf32>) -> tensor<1x2x2x3xf32>
    %2 = "tf.Relu6"(%1) {device = ""} : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
    %3 = "tf.Identity"(%2) {device = ""} : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
    return %3 : tensor<1x2x2x3xf32>
  }
}

// CHECK-LABEL: func @depthwise_conv

  %1 = mhlo.multiply %0, %cst1 : tensor<1x1x2x4xf32>
  // CHECK:      return %[[RES]] : tensor<1x1x2x4xf32>
  func.return %1 : tensor<1x1x2x4xf32>
}

// CHECK-LABEL: @foldBroadcastInDimBeforeMulOp_bcast_dim_2D_float
func.func @foldBroadcastInDimBeforeMulOp_bcast_dim_2D_float() -> (tensor<1x2x2x3xf32>) {

  %1 = "tf.Sub"(%0, %cst_0) {data_format = "NHWC"} : (tensor<1x3x2x2xf32>, tensor<1x1x1x2xf32>) -> tensor<1x3x2x2xf32>
  %2 = "tf.Mul"(%1, %cst_1) : (tensor<1x3x2x2xf32>, tensor<1x1x1x2xf32>) -> tensor<1x3x2x2xf32>
  func.return %2 : tensor<1x3x2x2xf32>
}

// CHECK: func @fuse_conv2d_with_sub_and_mul

  %1 = "tf.Maximum"(%0, %cst_0) : (tensor<1x3x4x2xf32>, tensor<f32>) -> tensor<1x3x4x2xf32>
  %2 = "tf.Minimum"(%1, %cst_1) : (tensor<1x3x4x2xf32>, tensor<f32>) -> tensor<1x3x4x2xf32>
  func.return %2 : tensor<1x3x4x2xf32>
// CHECK-DAG: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<{{.*}}> : tensor<1x1x3x2xf32>}> : () -> tensor<1x1x3x2xf32>

func.func @dont_move_transpose_different_ranks(%arg0:tensor<1x1x2x3xf32>, %arg1:tensor<2x3xf32>) -> tensor<1x2x1x3xf32> {
  %cst = "tf.Const"() {value = dense<[0, 2, 1, 3]> : tensor<4xi32>} : () -> tensor<4xi32>
  %0 = "tf.AddV2"(%arg0, %arg1) {device = ""} : (tensor<1x1x2x3xf32>, tensor<2x3xf32>) -> tensor<1x1x2x3xf32>
  %1 = "tf.Transpose"(%0, %cst) {device = ""} : (tensor<1x1x2x3xf32>, tensor<4xi32>) -> tensor<1x2x1x3xf32>

  func.return %1 : tensor<1x2x1x3xf32>

    %2 = "quantfork.stats"(%1) {layerStats = dense<[3.50919247, 6.000000e+00]> : tensor<2xf32>} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32>
    %3 = "tf.Identity"(%2) {device = ""} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32>
    %4 = "quantfork.stats"(%3) {layerStats = dense<[3.50919247, 6.000000e+00]> : tensor<2xf32>} : (tensor<1x3x4x2xf32>) -> tensor<1x3x4x2xf32>
    func.return %4 : tensor<1x3x4x2xf32>

    %13 = "tf.Mul"(%11, %12) : (tensor<1xf32>, tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
    %14 = "tf.Round"(%13) : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xf32>
    %15 = "tf.Cast"(%14) {Truncate = false} : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3xi32>
    %16 = "tf.AddV2"(%15, %arg10) : (tensor<1x2x2x3xi32>, tensor<i32>) -> tensor<1x2x2x3xi32>

}

func.func @cast(%arg0: tensor<1x2x2x5xi32>) -> tensor<1x2x2x5xf32> {
  %0 = "tf.Cast"(%arg0) : (tensor<1x2x2x5xi32>) -> tensor<1x2x2x5xf32>
  func.return %0 : tensor<1x2x2x5xf32>

  // CHECK-LABEL: cast
  // CHECK: "tfl.cast"(%arg0) : (tensor<1x2x2x5xi32>) -> tensor<1x2x2x5xf32>
}

func.func @castFloat32ToI16(%arg0: tensor<1x2x2x5xf32>) -> tensor<1x2x2x5xi16> {

// CHECK: stablehlo.maximum
// CHECK: (tensor<1x5x5x4xf32>, tensor<f32>) -> tensor<1x2x2x4xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %[[REDUCE_WINDOW_MAX]], dims = [0, 3, 1, 2] : (tensor<1x2x2x4xf32>) -> tensor<1x4x2x2xf32>
// CHECK: return %[[TRANSPOSE_1]]

// -----

// Tests that a `maximum(add(convolution(%activation, %weight), %bias), %zero)`

// CHECK: return %4 : tensor<1x1x2x2xf32>

func.func @avg_pool2d_6(%arg0: tensor<1x1x1x7xf32>) -> (tensor<1x1x1x2xf32>) {