func.func @unfuse_batch_norm(
    %x: tensor<4x256xf32>, %scale: tensor<256xf32>, %offset: tensor<256xf32>,
    %mean: tensor<256xf32>, %variance: tensor<256xf32>)
    -> (tensor<4x256xf32>) {
  // CHECK-DAG: %[[EPS_BCAST:.+]] = mhlo.constant dense<1.001000e-05> : tensor<256xf32>
  // CHECK-DAG: %[[VARIANCE_EPS:.+]] = mhlo.add %[[VARIANCE]], %[[EPS_BCAST]] : tensor<256xf32>

  %2 = "tfl.batch_matmul"(%1, %arg2) {adj_x = true, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<4x8xf32>, tensor<4x256xf32>) -> tensor<8x256xf32>
  func.return %2 : tensor<8x256xf32>
  // CHECK: return %[[RES1]] : tensor<8x256xf32>
}

// CHECK-LABEL: Batchmatmul2Fullyconnected
// CHECK-NOT: "tfl.batch_matmul"

       {
         data_format = "NCHW",
         epsilon = 1.001000e-05 : f32
       } : (tensor<?x256x56x56xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>)
        -> (tensor<?x256x56x56xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>, tensor<256xf32>, tensor<*xf32>)

  // CHECK: %[[BATCH_NORM1:[_a-z0-9]*]], {{.*}} = "tf.FusedBatchNormV3"
  // CHECK-SAME: %[[CONV1]]

func.func @main(tensor<4x5xbf16>) -> tensor<4x5xbf16> {
^bb0(%arg0: tensor<4x5xbf16>):
  // CHECK-LABEL: @main
  // CHECK:  (tensor<4x5xbf16>) -> tensor<4x5xf32>
  // CHECK-NEXT:  (tensor<4x5xf32>) -> tensor<4x5xbf16>
  %0 = "tfl.cast" (%arg0) : (tensor<4x5xbf16>) -> tensor<4x5xf32> loc("cast1")
  %1 = "tfl.cast" (%0) : (tensor<4x5xf32>) -> tensor<4x5xbf16> loc("cast2")
  func.return %1 : tensor<4x5xbf16>

  func.return %y : tensor<1x64x28x28xf32>
}

// CHECK-LABEL: bias_add_nchw
func.func @bias_add_nchw(%arg0: tensor<1x256x150x150xf32>, %arg1: tensor<256xf32>) -> tensor<1x256x150x150xf32> {
  // CHECK: (%[[ARG0:.*]]: tensor<1x256x150x150xf32>, %[[ARG1:.*]]: tensor<256xf32>)
  // CHECK: %[[CST:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}>
  // CHECK: %[[R0:.*]] = "tf.Transpose"(%[[ARG0]], %[[CST]])

// CHECK-NEXT:        buffer: 4
// CHECK-NEXT:      } ],
// CHECK-NEXT:      signature_defs: [  ]
// CHECK-NEXT:    }

  %0 = "tfl.cast" (%arg0) : (tensor<4x5xbf16>) -> tensor<4x5xf32> loc("cast1")
  %1 = "tfl.cast" (%0) : (tensor<4x5xf32>) -> tensor<4x5xbf16> loc("cast2")
  func.return %1 : tensor<4x5xbf16>