}

    int feature_group_count = conv_op.getFeatureGroupCount();
    // For depthwise and group convolutions, feature_group_count != 1
    if (feature_group_count != 1) {
      // Depthwise or Group convolution is not supported yet.
      return rewriter.notifyMatchFailure(
          conv_op, "group or depthwise convolution is not supported");
    }

    // Constructs strides array from lhs_dilation.

  return new_filter_constant_value_attr;
}

// Checks if the given convolution op is depthwise.
bool IsDepthwiseConvolution(stablehlo::ConvolutionOp op) {
  // `feature_group_count` controls how the input channel dimension is
  // split.
  // A value bigger than one signals depthwise convolution behavior.
  return op.getFeatureGroupCount() > 1;
}

  if (elements_depth == 1) {
    return true;
  }

  // In TFLite Conv2D uses OHWI format for filter, and 1HWO for Depthwise Conv.
  // For conv:
  // Check if last dimension in filter equals the first dimension
  // For depthwise conv:
  // Check if the first in filter dimension equals the first dimension.
  if (filter_shape.empty() ||
      (is_depthwise ? filter_shape.back() != elements_depth

^bb0(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<3x3x3x4xf32>) :
  %0 = "tf.DepthwiseConv2dNative"(%arg0, %arg1) {device = "", name = "MobilenetV2/expanded_conv/depthwise/depthwise", T = "tfdtype$DT_FLOAT", data_format = "NHWC", dilations = [1, 1, 1, 1], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<256x32x32x3xf32>, tensor<3x3x3x4xf32>) -> tensor<256x30x30x12xf32>

                shape=[
                    tensor_shape_pb2.TensorShapeProto(
                        dim=[
                            tensor_shape_pb2.TensorShapeProto.Dim(
                                # Depthwise conv is reshaped to [H,W,1,CxM].
                                size=filter_shape[quantized_axis]
                                * filter_shape[2]
                            )
                        ]

  );

  let results = (outs TFL_TensorOf<[F32, I32, I64]>:$output);

  let hasOptions = 1;
}

def TFL_DepthwiseConv2DOp :
    TFL_ConvOp<"depthwise_conv_2d", "Depthwise-separable convolution", 3,
                [DeclareOpInterfaceMethods<TFL_ArithmeticCount>,
                 DynamicRangeQuantizedOpInterface]> {
  let arguments = (
    ins TFL_TensorOf<[F32, QI8, QUI8, QI16]>:$input,