inline constexpr StringRef kQuantizationMethodAttr = "_quantization_method";

// Permutation from the NHWC tensor format to NCHW. This is an inverse
// permutation of `kNchwToNhwcPermutation`.
inline constexpr std::array<int64_t, 4> kNhwcToNchwPermutation = {0, 3, 1, 2};

// Permutation from the NCHW tensor format to NHWC. This is an inverse
// permutation of `kNchwToNhwcPermutation`.

  // Also FusedBatchNorm in training mode prefers NCHW data format. Check if all
  // users can efficiently use NHWC data format?
  if (one_by_one && trivial_strides && trivial_dilations) {
    return "NHWC";
  }

  // If filter spatial dimensions are unknown or not 1x1 we prefer NCHW, because
  // it's the fastest option on NVIDIA GPUs with cuDNN library support.
  return "NCHW";
}

ShapedType GetNhwcReturnTypeFromNchw(Operation* old_op) {
  auto composite_result_shape =
      mlir::cast<ShapedType>(old_op->getResults().front().getType()).getShape();
  std::array<int64_t, 4> output_shape;
  // NHWC <- NCHW
  output_shape[0] = composite_result_shape[0];
  output_shape[1] = composite_result_shape[2];
  output_shape[2] = composite_result_shape[3];
  output_shape[3] = composite_result_shape[1];

// through a StableHLO <-> MHLO roundtrip to utilize the MHLOQuantToInt pass.
void AddStablehloQuantToIntPasses(OpPassManager& pm);

// Processes tensors with NCHW format (== (batch, channel, height, weight)) by
// converting them to NHWC formats along with extra optimizations such as
// constant folding the transpose->convolution pattern. This is useful when

// RUN: stablehlo-quant-opt %s -stablehlo-nchw-convolution-to-nhwc \
// RUN:   -split-input-file -verify-diagnostics | FileCheck %s

// Tests that `stablehlo.transpose` ops are inserted for each of input, filter,
// and output.
// Output dimension numbers =  [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]

// CHECK-LABEL: nchw_conv
// CHECK-SAME: %[[ARG:.+]]: tensor<1x8x4x4xf32>
func.func @nchw_conv(%arg0: tensor<1x8x4x4xf32>) -> tensor<1x8x4x4xf32> {

      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>

def TF_ConvnetDataFormatAttr : StringBasedAttr<
    CPred<"$_self.cast<StringAttr>().getValue() == \"NHWC\" || " #
          "$_self.cast<StringAttr>().getValue() == \"NCHW\"">,
    "'NHWC' or 'NCHW' convnet data format">;

//===----------------------------------------------------------------------===//
// Type attributes

// A derived attribute that returns the size of `idx`-th ODS-declared variadic

# 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

node {
  name: "input"
  op: "Placeholder"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {