%3 = "tf.BatchMatMulV2"(%input, %2) {
      attr_map = "adj_x:0,adj_y:1"
    } : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
    func.return %3 : tensor<*xf32>
  }

  // DepthwiseConv2D with float computation
  func.func private @internal_depthwise_conv2d_fn(
                         %input : tensor<*xf32>, %filter : tensor<*xi8>) -> tensor<*xf32> {

          (TFL_DequantizeOp
              (TFL_QuantizeOp (TF_ReshapeOp $input, $shape),
              (UpdateShapeWithAxis<3> $qtype, $old_value))),
  [(UsedBy<"DepthwiseConv2D"> $old_value),
   (CanUpdateShapeWithAxis<3> $qtype, $old_value)],
  [], (addBenefit 10)>;

// The axis is set to 3, because this transpose is from the legalization of

  );

  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 = (

                         %input : tensor<*xf32>, %weight : tensor<*x!tf_type.qint8>,
                         %weight_scale : tensor<*xf32>, %weight_zp : tensor<*xi32>) -> tensor<*xf32>
      attributes {tf_quant.quantized_ops = ["DepthwiseConv2D"]} {

    %out = "tf.UniformQuantizedConvolutionHybrid"(%input, %weight,
                           %weight_scale, %weight_zp) {
        Tlhs = "tfdtype$DT_FLOAT",
        Trhs = "tfdtype$DT_QINT8",

    StringRef function_name =
        mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue();
    if (!function_name.starts_with("composite_")) {
      return spec;
    }
    if (function_name.contains("depthwise_conv2d")) {
      spec->coeff_op_quant_dim[1] = 3;
      if (function_name.contains("with_bias")) {
        spec->biases_params[2] = {{0, 1},
                                  quant::GetUniformQuantizedTypeForBias};
      }

      attr_map = "strides:0,use_cudnn_on_gpu:1,padding:2,explicit_paddings:3,dilations:4"
    } : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
    func.return %5 : tensor<*xi32>
  }

  // DepthwiseConv2D with float computation
  func.func private @internal_depthwise_conv2d_fn(
                         %input : tensor<*xi8>, %filter : tensor<*xi8>,

   (IsConstTensor $filter),
   (IsInt32ElementType $conv),
   (HasStaticShapeConstraint $filter),
   (HasStaticShapeAtDimsConstraint<"3"> $input)],
  [], (addBenefit 10)>;

// Converts inlined DepthwiseConv2D pattern to TF XlaConvV2 op. This pattern
// doesn't support non-constant weights.
def ConvertTFDepthwiseConv2DToXLAConvOp : Pat<
  (TF_CastOp:$conv
    (TF_DepthwiseConv2dNativeOp
      (TF_CastOp:$cast_input

    // than function name.
    if (!function_name.starts_with("composite_")) {
      return failure();
    }

    if (function_name.contains("depthwise_conv2d")) {
      // Uniform Quantized op requires weights of tf.DepthwiseConv2dNative to
      // be transformed from [H,W,C,M] to [H,W,1,CxM] where
      // H=height,W=width,C=channel,M=multiplier. Therefore, a reshape op is

      attr_map = "strides:0,use_cudnn_on_gpu:1,padding:2,explicit_paddings:3,dilations:4"
    } : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>
    func.return %5 : tensor<*xi32>
  }

  // DepthwiseConv2D with (simulated) int32 accumulation.
  func.func private @internal_depthwise_conv2d_fn(
                         %input : tensor<*xi8>, %filter : tensor<*xi8>,

// CHECK: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[DEQUANTIZED]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1]}> {attr_map = "0:strides,1:padding,2:explicit_paddings,3:dilations", device = ""} : (tensor<1x3x4x512xf32>, tensor<2x3x3x512xf32>) -> tensor<*xf32>
// CHECK: return %[[DEPTHWISE_CONV2D]] : tensor<*xf32>