%cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32>
  %cst_1 = "tf.Const"() {value = dense<3.000000e+00> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
  %0 = "tf.Conv2D"(%arg0, %cst_1) {
    data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [],
    padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true

  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"}
  ]

    } else if (function_name.contains("conv2d")) {
      // For Conv2D, the channel dimension must be static to calculate the
      // feature group count.
      if (!HasStaticShapeAtDims(call_op->getOperand(0), /*dims=*/3)) {
        return absl::InternalError(
            "The channel dimension of Conv2D is required to be static.");
      }
    } else if (function_name.contains("conv3d")) {

// CHECK: %[[CONV_1:.*]] = "tf.Conv2D"(%[[GATHER]], %[[DEQUANTIZED_1]]) <{data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> {device = ""} : (tensor<1x3x4x3xf32>, tensor<2x3x3x1024xf32>) -> tensor<1x3x2x1024xf32>

  %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({{.*}})"))
}

      auto conv_result = rewriter.create<mhlo::ConvolutionOp>(
          conv_op.getLoc(), new_output_type, sliced_input, sliced_kernel,
          conv_op.getWindowStridesAttr(), conv_op.getPaddingAttr(),
          conv_op.getLhsDilationAttr(), conv_op.getRhsDilationAttr(),
          conv_op.getWindowReversalAttr(), conv_op.getDimensionNumbers(), 1, 1,
          conv_op.getPrecisionConfigAttr());

      conv_results.push_back(conv_result);

// CHECK: %[[CONV2D:.*]] = "tf.Conv2D"(%arg0, %[[CONST]]) <{data_format = "NHWC", dilations = [1, 1, 2, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 2, 1], use_cudnn_on_gpu = true}> : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x2x2xf32>

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

    }
  }

  // Handle Conv2D input, stride and filter.
  HandleConv2DInput(conv2d, block_size);
  HandleConv2DStride(conv2d);
  HandleConv2DFilter(conv2d, block_size);

  // Book keeping new filter shape for backprop filter rewrite.
  // Filter shape is defined in HandleConv2DFilter, thus it is RankedTensorType.
  filter_shape =
      mlir::cast<RankedTensorType>(conv2d.getFilter().getType()).getShape();

  %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]]