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

  LogicalResult matchAndRewrite(mhlo::MulOp mul_op,
                                PatternRewriter &rewriter) const override {
    // Variables for capturing values and attributes used while creating ops.
    mhlo::ConvolutionOp conv_op;
    Operation *bcast_or_const_op;
    shape::ShapeOfOp shape_of_op;
    mhlo::ConstantOp filter;
    mhlo::ConstantOp multiplier;
    mlir::ElementsAttr filter_value, mul_value;

include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mhlo/IR/hlo_ops.td"
include "stablehlo/dialect/ChloOps.td"

def IsDefinedByConvOrDotGeneralOp : Constraint<Or<[
  CPred<"$0.getDefiningOp<mhlo::ConvolutionOp>()">,
  CPred<"$0.getDefiningOp<mhlo::DotGeneralOp>()">]>>;

def IsNull : Constraint<CPred<"!$0">>;

// This pattern optimizes:
//   conv/dot_general + a + b -> conv/dot_general + (a + b)