// 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

    %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

                         %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",

      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

      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>,

    // 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
      // feature group count.

//   ...${key2}...
// }
// ```
// The above template with generate two functions by substituting `key1` and
// `key2` with given values.

module {

  for main_op in ["Conv2D", "DepthwiseConv2D", "MatMul"] {
    parameters[
      {"quantized_ops": ["${main_op}", "BiasAdd"], "act_func": "internal_requantize_no_activation_fn", "output_type": "!tf_type.qint8"},