# CHECK: fake/user/code/file_C.py:27:1: error: 'tf.Conv2D' op attribute 'data_format' failed to satisfy constraint: 'NHWC' or 'NCHW' convnet data format
# CHECK: fake/user/code/file_D.py:28:1: note: called from
# CHECK: fake/user/code/file_E.py:29:1: note: called from
# CHECK: fake/user/code/file_F.py:30:1: note: called from

node {
  name: "input"

  auto var = ops::OnesLike(scope_, scale);

  if (!channel_first) {
    attrs.data_format_ = (shape.dims() == 5) ? "NDHWC" : "NHWC";
  } else {
    attrs.data_format_ = (shape.dims() == 5) ? "NCDHW" : "NCHW";
  }

  auto y = FusedBatchNormV3(scope_, x, scale, offset, mean, var, attrs);
  RunTest({x, scale, offset}, {shape, scale_shape, scale_shape}, {y.y},
          {shape});
}

INSTANTIATE_TEST_SUITE_P(

   // Unsupported data format
   %1 = "tf.Conv2D"(%arg2, %arg1) {T = "tfdtype$DT_FLOAT", data_format = "NCHW", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<256x3x32x32xf32>, tensor<3x3x3x16xf32>) -> tensor<256x16x32x32xf32>
   // OK

  // OK
  %0 = "tf.AvgPool"(%arg0) {T = "tfdtype$DT_FLOAT", data_format = "NCHW", ksize = [1, 1, 3, 6], padding = "VALID", strides = [1, 1, 3, 1]} : (tensor<1x16x6x6xf32>) -> tensor<1x16x1x1xf32>
  // Unsupported ksize
  %1 = "tf.AvgPool"(%arg0) {T = "tfdtype$DT_FLOAT", data_format = "NCHW", ksize = [3, 1, 3, 6], padding = "VALID", strides = [1, 1, 3, 1]} : (tensor<1x16x6x6xf32>) -> tensor<1x16x1x1xf32>

//===----------------------------------------------------------------------===//
def IsDataFormatNHWC : ConstantAttr<TF_ConvnetDataFormatAttr, "\"NHWC\"">;
def IsDataFormatNCHW : ConstantAttr<TF_ConvnetDataFormatAttr, "\"NCHW\"">;

class I32VectorElementsAttr<int len> : ElementsAttrBase<
  CPred<"$_self.isa<DenseIntElementsAttr>() &&"
      "$_self.cast<DenseIntElementsAttr>().getType()."
      "getElementType().isSignlessInteger(32)">,

  %0:6 = "tf.FusedBatchNormV3"(%arg0, %arg1, %arg2, %arg3, %arg4) {T = "tfdtype$DT_FLOAT", data_format = "NCHW", epsilon = 0.001 : f32, exponential_avg_factor = 1.0 : f32, is_training = true} : (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>,...

  %0 = "tf.BiasAddGrad"(%arg0) {data_format = "NCHW"} : (tensor<2x3x4x5xf32>) -> tensor<3xf32>
  func.return %0 : tensor<3xf32>
}

// CHECK-LABEL: func @BiasAddGrad_dynamic
func.func @BiasAddGrad_dynamic(%arg0: tensor<?x?x?x?xf32>) -> tensor<?xf32> {
  // CHECK: tf.Sum
  %0 = "tf.BiasAddGrad"(%arg0) {data_format = "NCHW"} : (tensor<?x?x?x?xf32>) -> tensor<?xf32>
  func.return %0 : tensor<?xf32>
}

}

// -----

func.func @testBiasAdd(%arg0: tensor<2x3xf32>, %arg1: tensor<3xf32>) -> tensor<2x3xf32> {
  // expected-error @+1 {{requires value operand to have rank at least three with `NCHW` data format}}
  %0 = "tf.BiasAdd"(%arg0, %arg1) {data_format = "NCHW"} : (tensor<2x3xf32>, tensor<3xf32>) -> tensor<2x3xf32>
  func.return %0 : tensor<2x3xf32>
}

// -----

```
With `src_format` of `NHWC`, `dst_format` of `NCHW`, and input:
```
[1, 2]
```
, the output will be:
```
[1, 2]
```
  }];

  let arguments = (ins
    Arg<TF_I32OrI64Tensor, [{Tensor of rank 1 or 2 in source data format.}]>:$x,

    DefaultValuedOptionalAttr<StrAttr, "\"NHWC\"">:$src_format,
    DefaultValuedOptionalAttr<StrAttr, "\"NCHW\"">:$dst_format
  );

  bool is_pool2d = rank == 4;
  if (verify_batch_channel_dims(0, rank - 1)) {
    *data_format = is_pool2d ? "NHWC" : "NDHWC";
  } else if (verify_batch_channel_dims(0, 1)) {
    *data_format = is_pool2d ? "NCHW" : "NCDHW";
  } else {
    return false;
  }

  if (rw.getWindowDilations().has_value() &&
      !(rw.getWindowDilations()->isSplat() &&
        rw.getWindowDilations()->getSplatValue<APInt>() == 1))