ConstBoolAttrFalse,
      ConstBoolAttrTrue),
    [(IsNhwcLayoutOp $attrs)]>;

// pattern to lower a stablehlo.composite of `jax` and `pytorch` image resize fuctions
// in `nearest`mode and with NCHW inputs to a tflite.resize_nearest_neighbor op.
// TODO(b/343278954): Move the creation of transposes to a separate prepare pass
// to avoid creating multiple pattern-rewrite rules for the same composite op.

class NchwConvolutionToNhwcPass
    : public impl::NchwConvolutionToNhwcPassBase<NchwConvolutionToNhwcPass> {
 private:
  void runOnOperation() override;
};

// Rewrites NCHW convolution to NHWC.
// * Src dimension numbers: [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1]
// * Dst dimension numbers: [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
class RewriteNchwConvolutionToNhwc

// CHECK: return %[[CUSTOM_AGGREGATOR_1]] : tensor<1x3xf32>
// CHECK: }
// CHECK: }

// -----

// Tests that `stablehlo.convolution` with NCHW format is converted to NHWC.

func.func @main(%arg0: tensor<1x8x4x4xf32>) -> tensor<1x8x4x4xf32> {
  %0 = stablehlo.constant() {value = dense<3.000000e+00> : tensor<8x8x3x3xf32>} : () -> tensor<8x8x3x3xf32>

// CHECK: stablehlo.reduce_window

// -----

// [No change] Tests that the transpose of the input of
// `stablehlo.reduce_window` is NOT deferred to the result, when the transpose
// isn't `[0, 3, 1, 2]` (i.e. NCHW->NHWC).

// CHECK-LABEL: reduce_window_max_activation_transpose_with_padding
func.func @reduce_window_max_activation_transpose_with_padding(%arg0: tensor<16x16x4x1xf32>) -> tensor<1x4x8x8xf32> {

   // 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>,...

```
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))