// For models with NCHW convolution format. This pass is required because
  // downstream pipeline handles NHWC convolution better for most cases.
  pm.addNestedPass<func::FuncOp>(createNchwConvolutionToNhwcPass());

  // Folds `stablehlo.constant`->`stablehlo.transpose` patterns, which is often
  // generated as by-products after optimizing dimension numbers (e.g.
  // NCHW->NHWC convolution conversion).

  // CHECK: %[[QUANT_CST:.*]] = stablehlo.uniform_quantize %[[CST]]
  // CHECK: %[[QUANT_ARG_0:.*]] = stablehlo.uniform_quantize %[[ARG_0]]
  // CHECK: %[[CONV:.*]] = stablehlo.convolution(%[[QUANT_ARG_0]], %[[QUANT_CST]])
  // CHECK-NOT: stablehlo.uniform_quantize
  // CHECK: %[[DEQUANT:.*]] = stablehlo.uniform_dequantize %[[CONV]]
  // CHECK: return %[[DEQUANT]]

// CHECK-SAME: (tensor<1x2xf32>, tensor<2x3x!quant.uniform<i8<-127:127>:f32, 0.0023622048182750312>>) -> tensor<1x3xf32>
// CHECK: return %[[DOT]]

// -----

// Test that per-tensor weight-only quantized convolution op is produced when
// empty `weight_only_ptq` is provided.

module attributes {tf_saved_model.semantics} {

  QuantizationSpec spec{};
  if (method_case != Method::kStaticRangePtq) {
    return spec;
  }

  // Matches all convolution quantizable unit family.
  spec.mutable_matcher()->mutable_function_name()->set_regex(
      "composite_conv.*");

  // Enable per-channel quantization for convolution weights.
  QuantizedType conv_weight_quantized_type{};

  // Assumes NHWC format, specifying the channel dimension (3) as the

  // NOFOLD: %[[PAD:[0-9]*]] = "tf.Pad"(%[[TRANSPOSE]], %[[PADDING]])

  // ------------------------------------------------------------------------ //
  // Convolution layer #0.
  // ------------------------------------------------------------------------ //
  %5 = "tf.Conv2D"(%4, %arg3)
        {
          data_format = "NCHW",
          dilations = [1, 1, 1, 1],

using ::stablehlo::quantization::QuantizedType;
using ::stablehlo::quantization::WeightOnlyPtq;

// Inserts quantization parameters of weights for weight-only quantization and
// dynamic range quantization of `stablehlo.convolution` and
// `stablehlo.dot_general`.
class InsertWeightParamPass
    : public impl::InsertWeightParamPassBase<InsertWeightParamPass> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertWeightParamPass)

      type: DT_FLOAT
    }
  }
  attr {
    key: "_class"
    value {
      list {
        s: "loc:@conv_net_2d/conv_2d_0/w"
      }
    }
  }
}
node {
  name: "conv_net_2d_1/conv_2d_0/convolution"
  op: "Conv2D"
  input: "input"
  input: "conv_net_2d/conv_2d_0/w/read"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "data_format"
    value {

  let dependentDialects = ["mlir::stablehlo::StablehloDialect",];
}

def NchwConvolutionToNhwcPass : Pass<"stablehlo-nchw-convolution-to-nhwc", "mlir::func::FuncOp"> {
  let summary = "Converts stablehlo.convolution op of NCHW format to -> NHWC.";
  let description = [{
    Matches `ConvolutionOp`s with NCHW format and converts it to NHWC

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

namespace mlir {
namespace TFL {

// A dilated convolution can be emulated with a regular convolution by chaining
// SpaceToBatch and BatchToSpace ops before and after it:
//
//     SpaceToBatchND -> Conv2D -> BatchToSpaceND
//
// This method was common before Conv2D fully supported dilated convolution in
// TensorFlow. This transformation detects this "emulation", and replaces it
// with a true dilated convolution, eliminating the SpaceToBatch and