if (!input_shape) {
    return op->emitError(
        "Only input with known shape is supported for Uniform Quantized "
        "opset.");
  }

  if (op->getParentOfType<func::FuncOp>().getName().contains("depthwise_")) {
    feature_group_cnt = input_shape.getDimSize(3);
  }

  attrs.push_back(rewriter.getNamedAttr(
      feature_group_cnt_attr, rewriter.getI64IntegerAttr(feature_group_cnt)));

// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op | FileCheck %s

module {
  // For UniformQuantized depthwise convolution, tensor shape should have
  // transformed from [H,W,C,M] to [H,W,1,CxM],
  func.func @depthwise_conv(%arg0: tensor<1x3x4x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<6xf32>} : () -> tensor<6xf32>

module {
  // For XLA weight-only per-channel depthwise convolution, tensor shape should have
  // transformed from [H,W,C,M] to [H,W,1,CxM],
  func.func @depthwise_conv(%arg0: tensor<1x3x4x3xf32>) -> (tensor<*xf32>) {

    }

    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.

                "non-broadcastable operands";
      });
    }
    filter_value = filter.getValue();
    mul_value = multiplier.getValue();
    // In MHLO, Conv filter is in HWIO format, Depthwise conv filter is in HW1O
    // format and backprop input conv filter is in HWOI format.
    // Only fuses multiplier if all dimensions other than the out channel
    // dimension are equal to 1.

  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;
}

    func.return %conv_out : tensor<*x!tf_type.qint32>
  }

  // Depthwise convolution. feature_group_count is set to 3rd dim of input shape.
  func.func private @internal_depthwise_conv2d_fn(
                         %input : tensor<*x!tf_type.qint8>, %filter : tensor<*x!tf_type.qint8>,

    return success();
  }

  OpSet target_opset_;
};

// To calculate per-channel scale and offset, weight of depthwise was reshaped
// to [H, W, 1, InxMul]. After scale and offset has been calculated, this
// pattern gets called and restores the weight of depthwise back
// into [H, W, In, Mul]
class RestoreWeightShapePattern
    : public OpRewritePattern<TF::PartitionedCallOp> {

      strides: Sequence[int] = (1, 2, 2, 1),
      dilations: Sequence[int] = (1, 1, 1, 1),
      padding: str = 'SAME',
  ):
    class DepthwiseConvModel(module.Module):
      """A simple model with a single depthwise conv2d, bias and relu."""

      def __init__(self):
        self.out_channel_size = filter_shape[2] * filter_shape[3]

        # This ensures filters will have different value range per out channel

  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