ConstBoolAttrFalse,
      ConstBoolAttrTrue),
    [(IsNhwcLayoutOp $attrs)]>;

// pattern to lower a stablehlo.composite of `jax` and `pytorch` image resize fuctions
// in `nearest`mode and with NCHW inputs to a tflite.resize_nearest_neighbor op.
// TODO(b/343278954): Move the creation of transposes to a separate prepare pass
// to avoid creating multiple pattern-rewrite rules for the same composite op.

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 %[[CUSTOM_AGGREGATOR_1]] : tensor<1x3xf32>
// 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>