// %7 = stablehlo.convolution(%4, %6)  // q1 * q2 (disguised in f32).
// %8 = stablehlo.reshape %2  // z1
// %9 = stablehlo.broadcast_in_dim %8
// %10 = stablehlo.convert %9  // i8 -> f32 cast trick for z1.
// %11 = stablehlo.convert %5  // i8 -> f32 cast trick for filter.
// %12 = stablehlo.convolution(%10, %11)  // q2 * z1
// %13 = stablehlo.subtract %7, %12  // q1 * q2 - q2 * z1

                          %out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x${output_type}>
    attributes {tf_quant.quantized_ops = ${quantized_ops}} {
      // Given the convolution takes 2 qint8 inputs and output a qint32.
      // The accumulation scale is (input_scale * filter_scale).
      // The accumulation zero point is 0.

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

          ]
      )
      def depthwise_conv(
          self, input_tensor: core.Tensor
      ) -> Mapping[str, core.Tensor]:
        """Performs a 2D depthwise convolution operation.

        Args:
          input_tensor: Input tensor to perform convolution on.

        Returns:
          A map of: output key -> output result.
        """
        scale = [1.0] * self.out_channel_size

    %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
    return %0 : tensor<1x3x4x2xf32>
  }
  // CHECK: func private @composite_conv_fn
  // CHECK: %[[CONV:.+]] = stablehlo.convolution
  // CHECK: return %[[CONV]]
}

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();

module {
  // For XLA weight-only per-channel 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>) {

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