// CHECK-SAME: %[[PAD]]
  // CHECK-SAME: data_format = "NHWC"
  // CHECK-SAME: strides = [1, 2, 2, 1]

  %6, %batch_mean, %batch_variance, %reserved_1, %reserved_2, %reserved_3 =
       "tf.FusedBatchNormV3"(%5, %arg1, %arg1, %arg1, %arg1)
       {
         data_format = "NCHW",
         epsilon = 1.001000e-05 : f32,
         is_training = false

//   %cst_2 = "tf.Const"() {value = dense<[0.3, 0.4, 0.3, 0.4, 0.3, 0.4, 0.3, 0.4]> : tensor<8xf32>} : () -> tensor<8xf32>

  %arg0: tensor<1x28x28x64xf32>,
  %arg1: tensor<64xf32>
) -> tensor<1x28x28x64xf32> {

  // CHECK: "tf.FusedBatchNormV3"
  // CHECK-SAME: (%[[X_TRANSPOSE:[0-9]*]], %arg1, %arg1, %arg1, %arg1)
  // CHECK-SAME: data_format = "NCHW"
  %y, %batch_mean, %batch_var, %reserve_1, %reserve_2, %reserve_3
    = "tf.FusedBatchNormV3"(%arg0, %arg1, %arg1, %arg1, %arg1)
       {
         data_format = "NHWC",
         epsilon = 1.001 : f32,

// CHECK-NEXT: %[[Y:[a-z0-9]*]], {{.*}}_FusedBatchNormEx
// CHECK-NEXT: return %[[Y]]

 public:
  void runOnOperation() final;
};

//   %y:6 = "tf.FusedBatchNormV3"(%x, %scale, %offset, %mean, %variance)
//   %0 = "tf.Relu"(%y#0)
// ->
//   %y:6 = "tf._FusedBatchNormEx"(%x, %scale, %offset, %mean, %variance)
//
// Or:
//   %y:6 = "tf.FusedBatchNormV3"(%x, %scale, %offset, %mean, %variance)
//   %0 = "tf.AddV2"(%y#0, %side_input)
//   %1 = "tf.Relu"(%0)
// ->

func.func @fused_batch_norm_v3_bf16(%arg0: tensor<8x8x8x8xbf16>, %arg1: tensor<8xf32>, %arg2: tensor<8xf32>, %arg3: tensor<8xf32>, %arg4: tensor<8xf32>) -> tensor<8x8x8x8xbf16> {

  %arg0: tensor<1x28x28x64xf32>,
  %arg1: tensor<64xf32>
) -> tensor<1x28x28x64xf32> {

  // CHECK: "tf.FusedBatchNormV3"
  // CHECK-SAME: (%[[X_TRANSPOSE:[0-9]*]], %arg1, %arg1, %arg1, %arg1)
  // CHECK-SAME: data_format = "NCHW"
  %y, %batch_mean, %batch_var, %reserve_1, %reserve_2, %reserve_3
    = "tf.FusedBatchNormV3"(%arg0, %arg1, %arg1, %arg1, %arg1)
       {
         data_format = "NHWC",
         epsilon = 1.001 : f32,

// CHECK:  %[[ADD2:.*]] = "tf.Add"(%[[MUL2]], %[[SUB]])
// CHECK:  %[[BATCHNORM1_a:[^,]+]], {{.*}} = "tf.FusedBatchNormV3"(%[[ADD2]], %[[CONSTANT1]], %[[ARG2]], %[[ARG3]], %[[ARG4]])
// CHECK:  %[[BATCHNORM1_b:[^,]+]], {{.*}} = "tf.FusedBatchNormV3"(%[[BATCHNORM1_a]], %[[CONSTANT1]], %[[ARG2]], %[[ARG3]], %[[ARG4]])
// CHECK:  "tf.FusedBatchNormV3"(%[[BATCHNORM1_b]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]])
}

  // CHECK: %[[ARG_PERM:.*]] = "tf.Const"()
  // CHECK-SAME: <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}>
  // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]])

  // CHECK: "tf.FusedBatchNormV3"
  // CHECK-SAME: (%[[ARG_TRANSPOSE]], %arg1, %arg1, %arg1, %arg1)
  // CHECK-SAME: data_format = "NHWC"
  // CHECK-SAME: (tensor<1x28x28x64xf32>, tensor<64xf32>,
  // CHECK-SAME: -> (tensor<1x28x28x64xf32>, tensor<64xf32>,

  %1 = "tf.Transpose"(%arg0, %0) : (tensor<1x64x112x112xf32>, tensor<4xi32>) -> tensor<1x112x112x64xf32>

  // Compute FusedBatchNormV3 in NHWC format
  %2, %batch_mean, %batch_var, %reserve_1, %reserve_2, %reserve_3
    = "tf.FusedBatchNormV3"(%1, %arg1, %arg1, %arg1, %arg1)
       {
         data_format = "NHWC",
         epsilon = 1.001 : f32,
         exponential_avg_factor = 1.0 : f32,