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

// -----

// Test that hybrid quantized convolution is produced when q/dq pair only exists
// for weight.

module attributes {tf_saved_model.semantics} {

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

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

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

    }

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

  pm.addNestedPass<func::FuncOp>(createFuseConvolutionPass());
  pm.addNestedPass<func::FuncOp>(createOptimizePass());
  // Conditionally enable below pass because this causes unfused convolutions
  // described in b/293149194. This problem is not replicated in
  // StableHLO Quantizer.
  if (add_fold_broadcast_pass) {
    pm.addNestedPass<func::FuncOp>(createFoldBroadcastPass());
  }