// 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
// downstream pipeline (e.g. XLA) is more optimized when accepting NHWC formats.
void AddProcessNchwTensorPasses(OpPassManager& pm);

// usually has the following pattern. In the example below,
// the input operand would be stablehlo.convolution op, and return value would
// be stablehlo.add op.
//
// ```
// %0 = stablehlo.constant dense<3>
// %1 = stablehlo.constant dense<4>
// %2 = stablehlo.constant dense<2>
// %3 = stablehlo.convolution(%%arg0, %%arg1) :
//          (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<?x3x4x2xf32>

    }

    // If all quantized or floating point then types are consistent.
    // Int is valid in combination with both quantized and floating point.
    // This occurs when doing qi16 convolution, as bias is passed as a
    // non-quantized int64
    if (allTypesFp || allTypesQuantizedOrInt) return failure();

    Location loc = op->getLoc();
    SmallVector<Value> newOperands;

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 %[[VAL_14]] : tensor<16x32x256xbf16>
// CHECK:         }
func.func @convert_conv1d(%arg0: tensor<16x32x256xbf16>, %arg1: tensor<1x256x256xbf16>) -> tensor<16x32x256xbf16> {
	%0 = "mhlo.convolution"(%arg0, %arg1) {
    batch_group_count = 1 : i64,
    dimension_numbers = #mhlo.conv<[b, 0, f]x[0, i, o]->[b, 0, f]>,
    feature_group_count = 1 : i64,
    lhs_dilation = dense<1> : tensor<1xi64>,

        'b4': tf.Variable(tf.zeros([num_classes])),
    }

  @tf.function
  def __call__(self, data):
    """The Model definition."""
    x = tf.reshape(data, [-1, 28, 28, 1])

    # 2D convolution, with 'SAME' padding (i.e. the output feature map has
    # the same size as the input).

    # NOTE: The data/x/input is always specified in floating point precision.
    # output shape: [-1, 28, 28, 32]

}

// CHECK-LABEL: func @transposeConv2D_1x1_f32
func.func @transposeConv2D_1x1_f32(%input: tensor<1x64x28x28xf32>, %filter: tensor<1x1x64x64xf32>) -> tensor<1x64x28x28xf32> {
  // 1x1 convolution can be computed as a GEMM in NHWC data format.
  // CHECK: "tf.Conv2D"(%[[INPUT_TRANSPOSE:[0-9]*]], %arg1)
  // CHECK-SAME: data_format = "NHWC"
  %0 = "tf.Conv2D"(%input, %filter)
       {
         data_format = "NCHW",

  // CHECK: %[[CONST:.*]] = mhlo.constant()
  // CHECK-SAME{LITERAL} value = dense<127> : tensor<2x3x3x2xi8>
  // CHECK-SAME: tensor<2x3x3x2x!quant.uniform<i8:f32, 1.000000e+00:3>>
  // CHECK: mhlo.convolution(%arg0, %[[CONST]])
  // CHECK-SAME{LITERAL}: dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
  // CHECK-SAME{LITERAL}: window = {stride = [1, 2], pad = [[0, 0], [0, 0]], lhs_dilate = [1, 1], rhs_dilate = [2, 2]}

      return failure();
    }
    // Check if the given op is based on grouped convolution.
    // Dim size zero will be verified by the tf.Conv2D operator verification.
    if (input_type.getDimSize(3) % filter_type.getDimSize(2) != 0) {
      return failure();
    }

    // TensorFlow convolution op only has two inputs, while the TFLite one has

class DeferActivationTransposeForAddOp : public OpRewritePattern<AddOp> {
 public:
  using OpRewritePattern<AddOp>::OpRewritePattern;

  LogicalResult match(AddOp op) const override {
    // Only supports the case for 2D convolution.
    const Value lhs = op.getOperand(0);
    if (!HasRankOf(lhs, /*rank=*/4)) return failure();

    const Value rhs = op.getOperand(1);
    Operation* rhs_op = rhs.getDefiningOp();