// CHECK-NOT: func private @internal_conv2d_fn
// CHECK-NOT: func private @internal_matmul_fn
// CHECK: func private @quantized_conv2d_with_bias_fn
// CHECK-SAME: tf_quant.quantized_ops = ["Conv2D", "BiasAdd"]
// CHECK: func private @quantized_conv2d_with_bias_and_relu_fn
// CHECK: func private @quantized_conv2d_with_bias_and_relu6_fn
// CHECK: func private @quantized_conv2d_fn

                                PatternRewriter& rewriter) const override;
};

// Ensure bias for conv2d op.
struct EnsureBiasForConv2d : public OpRewritePattern<TFL::Conv2DOp> {
  using OpRewritePattern<TFL::Conv2DOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TFL::Conv2DOp conv_op,
                                PatternRewriter& rewriter) const override;
};

// Pad slice to 4d.

  parameters[
    {"quantized_ops": ["MatMul"], "internal_func_name": "internal_matmul_fn"},
    {"quantized_ops": ["Conv2D"], "internal_func_name": "internal_conv2d_fn"},
    {"quantized_ops": ["DepthwiseConv2D"], "internal_func_name": "internal_depthwise_conv2d_fn"},
    {"quantized_ops": ["Conv3D"], "internal_func_name": "internal_conv3d_fn"},
    {"quantized_ops": ["BatchMatMul"], "internal_func_name": "internal_batch_matmul_fn"}
  ]

  %2 = "tf.Cast"(%1) {Truncate = false} : (tensor<1x3x2x2xbf16>) -> tensor<1x3x2x2xf32>
  %3 = "tf.IdentityN"(%2) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>
  return %3 : tensor<1x3x2x2xf32>
// CHECK: tf.Conv2D
// CHECK-SAME: loc(callsite("Model/conv2d@conv2d_with_valid_loc"("Conv2D") at "QuantizationUnit({{.*}})"))
}

    %1 = "tf.Conv2D"(%0, %cst) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}
        : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32> loc(fused["Conv2D:", "Model/conv2d"])
    %2 = "tf.IdentityN"(%1) {device = ""} : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2xf32>

  %4 = "tf.Transpose"(%2, %3) : (tensor<1x32x32x8xf32>, tensor<4xi32>) -> tensor<1x8x32x32xf32>

  // Check that Conv2D computed in NCHW format, and all redundant transpose
  // operations removed from the function.

  // CHECK: %[[CONV:[0-9]*]] = "tf.Conv2D"(%arg0, %arg1)
  // CHECK-SAME: data_format = "NCHW"
  // CHECK-SAME: -> tensor<1x8x32x32xf32>

  // CHECK: return %[[CONV]]

      if (function_name.contains("with_bias")) {
        spec->biases_params[2] = {{0, 1},
                                  quant::GetUniformQuantizedTypeForBias};
      }
    } else if (function_name.contains("conv2d")) {
      spec->coeff_op_quant_dim[1] = 3;
      if (function_name.contains("with_bias")) {
        spec->biases_params[2] = {{0, 1},
                                  quant::GetUniformQuantizedTypeForBias};
      }

    return %3 : tensor<?x?x?x2xf32>
  }

// CHECK-LABEL: func @conv_with_dynamic_shape
// The Conv2D should not be quantized since it has dynamic channel.
// CHECK: "tf.Conv2D"
// CHECK-SAME: (tensor<?x?x?x?xf32>, tensor<2x3x3x2xf32>) -> tensor<?x?x?x2xf32>

  // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi64>}>
  // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]])

  // CHECK: %[[CONV2D:[0-9]*]] = "tf.Conv2D"(%[[ARG_TRANSPOSE]], %arg1)
  // CHECK-SAME: data_format = "NCHW"
  // CHECK-SAME: dilations = [1, 4, 2, 3]
  // CHECK-SAME: explicit_paddings = [1, 2, 7, 8, 3, 4, 5, 6]
  // CHECK-SAME: padding = "EXPLICIT"

  // cuDNN prefers NCHW data format for spatial convolutions.
  // CHECK: "tf.Conv2D"(%[[INPUT_TRANSPOSE:[0-9]*]], %arg1)
  // CHECK-SAME: data_format = "NCHW"
  %0 = "tf.Conv2D"(%input, %filter)
       {
         data_format = "NHWC",
         padding = "VALID",
         strides = [1, 1, 1, 1]