} 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")) {

      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>

    }
  }

  // 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();

      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};
      }

  }
  func.func private @composite_conv2d_with_bias_and_relu6_fn_1(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>, %arg2: tensor<2xf32>) -> tensor<*xf32> attributes {tf_quant.composite_function} {

//
//
//   SpaceToBatchND -> Expand -> Conv2D -> Squeeze -> BatchToSpaceND -> BiasAdd
//
//   SpaceToBatchND -> Expand -> Conv2D -> Squeeze -> Pad -> BatchToSpaceND ->
//   BiasAdd
//
//   SpaceToBatchND -> Expand -> Conv2D -> Squeeze -> BiasAdd -> BatchToSpaceND
//
//   SpaceToBatchND -> Conv2D -> Pad -> BatchToSpaceND -> BiasAdd
//
//   SpaceToBatchND -> Conv2D -> BatchToSpaceND -> BiasAdd
//
//

        mul_op.getLoc(), conv_op.getType(), conv_op.getLhs(), new_filter,
        conv_op.getWindowStridesAttr(), conv_op.getPaddingAttr(),
        conv_op.getLhsDilationAttr(), conv_op.getRhsDilationAttr(),
        conv_op.getWindowReversalAttr(), conv_op.getDimensionNumbers(),
        conv_op.getFeatureGroupCount(), conv_op.getBatchGroupCount(),
        conv_op.getPrecisionConfigAttr());

      equation = "",
      attr_map = "equation:0"
    } : (tensor<*xi32>, tensor<*xi32>) -> tensor<*xi32>

    func.return %4 : tensor<*xi32>
  }

  for main_op in ["Conv2D", "DepthwiseConv2D", "MatMul", "Conv3D", "BatchMatMul", "Einsum"] {
    parameters[
      {"quantized_ops": ["${main_op}", "BiasAdd"], "act_func": "internal_requantize_no_activation_fn", "output_type": "i8"},

  %cst_0 = "tf.Const"() {value = dense<-1.000000e+00> : tensor<f32>} : () -> tensor<f32>
  %cst_1 = "tf.Const"() {value = dense<1.000000e+00> : tensor<f32>} : () -> tensor<f32>
  %0 = "tf.Conv2D"(%arg0, %cst) {data_format = "NHWC", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x3x4x3xf32>, tensor<1x1x3x2xf32>) -> tensor<1x3x4x2xf32>