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 {

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

      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 {

    : public PassWrapper<FuseMhloConvolutionPass, OperationPass<func::FuncOp>> {
 public:
  StringRef getArgument() const final { return "fuse-mhlo-convolution-pass"; }
  StringRef getDescription() const final {
    return "Fuses MHLO binary element-wise ops and convolution op";
  }

  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());

// RUN: tf-opt %s -tf-layout-assignment=force-data-format=NHWC -verify-diagnostics | FileCheck %s --dump-input=always

// IMPORTANT: Tensor shapes do not match convolution parameters (stride,
// dilations, etc...). This test only verifies that changing convolution data
// layout will update all the attributes.

// CHECK-LABEL: func @transposeConv2D
func.func @transposeConv2D(%input: tensor<1x3x32x32xf32>, %filter: tensor<1x1x3x8xf32>) -> tensor<1x8x7x6xf32> {

    %1 = "quantfork.stats"(%arg0) {layerStats = dense<[1.27501142, 2.824783]> : tensor<2xf32>} : (tensor<1x3x2x3xf32>) -> tensor<1x3x2x3xf32>
    %2 = stablehlo.convolution(%1, %0)
      dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f],
      window = {
        stride = [1, 1], pad = [[0, 0], [1, 1]],
        lhs_dilate = [1, 1],
        rhs_dilate = [1, 1]
      }
      {

// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op | FileCheck %s

module {
  // For UniformQuantized depthwise convolution, tensor shape should have
  // transformed from [H,W,C,M] to [H,W,1,CxM],
  func.func @depthwise_conv(%arg0: tensor<1x3x4x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<6xf32>} : () -> tensor<6xf32>

  %4 = "tf.ReadVariableOp"(%arg1) : (tensor<!tf_type.resource<tensor<3x3x1x5xf32>>>) -> tensor<3x3x1x5xf32>
  %5 = "tf.ReadVariableOp"(%arg3) : (tensor<!tf_type.resource<tensor<3920x10xf32>>>) -> tensor<3920x10xf32>

// IMPORTANT: In the following Conv2D tests tensor shapes do not match
// convolution parameters (stride, dilations, etc...). This test only verifies
// that changing convolution data layout will update all the attributes.

// CHECK-LABEL: func @transposeConv2D

std::unique_ptr<Pass> createUnfuseBatchNormPass();

// Creates a pass which constant folds broadcast_in_dim op conditionally.
std::unique_ptr<Pass> createFoldBroadcastPass();

// Creates a pass which fuses MHLO binary element-wise ops and convolution op.
std::unique_ptr<Pass> createFuseConvolutionPass();

// Creates a pass which applies various optimizations on MHLO IR.
std::unique_ptr<Pass> createOptimizePass();