// * Dst dimension numbers: [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
class RewriteNchwConvolutionToNhwc
    : public OpRewritePattern<mlir::stablehlo::ConvolutionOp> {
 public:
  using OpRewritePattern<mlir::stablehlo::ConvolutionOp>::OpRewritePattern;

  LogicalResult match(mlir::stablehlo::ConvolutionOp op) const override {
    // Handles 2D convolutions only.
    if (!HasRankOf(op.getOperand(0), /*rank=*/4) ||

//   * Not a depthwise convolution.
class RewriteQuantizedConvolutionOp
    : public OpRewritePattern<stablehlo::ConvolutionOp> {
 public:
  using OpRewritePattern<stablehlo::ConvolutionOp>::OpRewritePattern;
  LogicalResult match(stablehlo::ConvolutionOp op) const override {
    const bool has_i32_output =
        IsI32F32UniformQuantizedPerAxisType(GetElementType(op.getResult()));

  }
}

template <int num_spatial_dims>
class ConvertNdConvOp : public OpConversionPattern<mhlo::ConvolutionOp> {
 public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(
      mhlo::ConvolutionOp conv_op, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final {
    SetDefaultConvAttributes(conv_op, rewriter);

// ```
class ComposeUniformQuantizedConvolutionOp
    : public OpRewritePattern<stablehlo::ConvolutionOp> {
 public:
  using OpRewritePattern<stablehlo::ConvolutionOp>::OpRewritePattern;

  LogicalResult match(stablehlo::ConvolutionOp op) const final {
    // Verify operands' types.
    for (Type operand_type : op.getOperandTypes()) {
      if (Type element_type =

    if (!quantization_method.has_static_range_ptq()) {
      return failure();
    }
    return MatchGemmStyleOp<ConvolutionOp>(entry_func_op);
  }

  void rewrite(func::FuncOp entry_func_op, const Method& quantization_method,
               PatternRewriter& rewriter) const override {
    RewriteGemmStyleOp<ConvolutionOp>(
        entry_func_op, rewriter,
        enable_per_channel_quantized_weight_ &&

#include "tsl/platform/status_matchers.h"

namespace mlir::quant {
namespace {

using ::mlir::stablehlo::AddOp;
using ::mlir::stablehlo::ConstantOp;
using ::mlir::stablehlo::ConvolutionOp;
using ::mlir::stablehlo::DotGeneralOp;
using ::mlir::stablehlo::SubtractOp;
using ::testing::ElementsAreArray;
using ::testing::Eq;
using ::testing::IsEmpty;
using ::testing::IsNull;
using ::testing::NotNull;

    const SymbolTable symbol_table(module_op);
    func::FuncOp func = symbol_table.lookup<func::FuncOp>(function_name);

    if (function_name.contains("conv")) {
      return (*(func.getOps<mlir::stablehlo::ConvolutionOp>().begin()))
          .getDimensionNumbers()
          .getKernelOutputFeatureDimension();
    } else if (function_name.contains("dot_general")) {
      auto dot = *(func.getOps<mlir::stablehlo::DotGeneralOp>().begin());

}

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
    format by inserting `TransposeOp`s to input, filter, and output tensors.
    In terms of dimension numbers, this matches

    if (failed(converted_attrs_or)) {
      return failure();
    }
    SmallVector<Value, 2> operands{op.getLhs(), *rhs_or};
    rewriter.replaceOpWithNewOp<mhlo::ConvolutionOp>(op, op.getType(), operands,
                                                     *converted_attrs_or);
    return success();
  }
};

class ConvertUniformQuantizeOp

                              batch_group_count_attr, paddings_attr,
                              precision_config_attr};
    rewriter.replaceOpWithNewOp<ConvolutionOp>(op, op.getType(), operands,
                                               llvm::ArrayRef(attrs));
    return success();
  }
};

using ConvertConv2DOp = ConvertConvOp<TF::Conv2DOp, /*num_spatial_dims=*/2>;