Type result_type, Value input,
                                     Value filter, Value bias) const {
    // Compared to tfl.conv_2d, tfl.depthwise_conv_2d has an additional
    // 'depth_multiplier' attribute. However, tf.DepthwiseConv2dNative does not
    // have a corresponding 'depth_multiplier' attribute; the multiplier is the

      for (int k = 0; k < filter_shape[2]; ++k) {
        for (int l = 0; l < filter_shape[3]; ++l) {
          // [o, 0, 1, i] for `tfl.conv_2d` case`,
          // [i, 0, 1, o] for `tfl.depthwise_conv_2d` case.
          int old_idx = get_array_idx(filter_shape, i, j, k, l);
          int new_idx = is_depthwise
                            ? get_array_idx(new_filter_shape, k, i, j, l)

  %w = arith.constant dense<127.0> : tensor<32x3x3x3xf32>
  %b = arith.constant dense<0.0> : tensor<32x1xf32>
  // expected-error @+1 {{'tfl.depthwise_conv_2d' op operand #2 must be 1D tensor of any type values or none type, but got 'tensor<32x1xf32>'}}

      succeeded(applyPatternsAndFoldGreedily(test_func, std::move(patterns))));

  auto depthwise_conv_op =
      FindOperationOfType<TF::DepthwiseConv2dNativeOp>(test_func);
  EXPECT_THAT(depthwise_conv_op, NotNull());
  // The filter of the DepthwiseConv2dNativeOp is expected to be a constant.
  EXPECT_TRUE(isa<TF::ConstOp>(depthwise_conv_op.getFilter().getDefiningOp()));
}

TEST_F(ConstantFoldingTest, DepthwiseConvWeightNotFoldable) {

// CHECK-NEXT: %[[DEPTHWISE_CONV2D:.*]] = "tf.DepthwiseConv2dNative"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]}> {device = ""} : (tensor<*xf32>, tensor<2x3x3x1xf32>) -> tensor<?x?x?x3xf32>

// CHECK: func private @quantized_conv2d_with_relu6_fn
// CHECK: func private @quantized_depthwise_conv2d_with_bias_and_relu_float_output_fn
// CHECK-SAME: tf_quant.quantized_ops = ["DepthwiseConv2D", "BiasAdd", "Relu"]
// CHECK: func private @quantized_matmul_with_bias_fn
// CHECK: func private @quantized_matmul_with_bias_and_relu_fn
// CHECK: func private @quantized_matmul_with_bias_and_relu6_fn

// CHECK: func private @quantized_conv2d_fn
// CHECK-SAME: tf_quant.quantized_ops = ["Conv2D"]
// CHECK: func private @quantized_depthwise_conv2d_fn
// CHECK-SAME: tf_quant.quantized_ops = ["DepthwiseConv2D"]

// UQ-CHECK: func private @quantized_conv2d_fn
// UQ-CHECK: func private @quantized_depthwise_conv2d_fn

    // Disable quantization for the DepthwiseConv since it has no benefits in
    // the XLA opset.
    if (function_name.contains("depthwise_conv2d")) {
      return absl::InternalError(
          "DepthwiseConv2D doesn't get any benefit of quantization in XLA.");
    } else if (function_name.contains("conv2d")) {
      // For Conv2D, the channel dimension must be static to calculate the

    %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