%0 = "tf.Conv3D"(%arg0, %cst) {
    data_format = "NDHWC", device = "", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 2, 1, 1]
  } : (tensor<1x3x4x3x3xf32>, tensor<2x3x3x3x2xf32>) -> tensor<1x3x2x3x2xf32>
  %1 = "tf.Relu"(%0) {device = ""} : (tensor<1x3x2x3x2xf32>) -> tensor<1x3x2x3x2xf32>

  %2 = "tf.Conv3D"(%arg0, %cst) {

// CHECK: %[[conv3d:.*]] = "tfl.conv_3d"(%arg0, %[[w]], %[[const]]) <{dilation_d_factor = 1 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_d = 1 : i32, stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<?x28x28x28x8xf32>, tensor<3x3x3x8x16xf32>, none) -> tensor<?x26x26x26x16xf32>
// CHECK: %2 = "tfl.shape"(%[[conv3d]]) : (tensor<?x26x26x26x16xf32>) -> tensor<5xi64>

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

  if (!TFPaddingIsSameOrValid(op, &padding)) return failure();

  // TensorFlow Conv3D has no bias, optimization patterns will fuse Conv3D
  // with other ops can fill the bias.
  Value none = rewriter.create<TFL::NoValueOp>(
      op->getLoc(), rewriter.getNoneType(), rewriter.getUnitAttr());

  rewriter.replaceOpWithNewOp<TFL::Conv3DOp>(
      op, tf_op.getType(), tf_op.getInput(), tf_op.getFilter(),

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

func.func @lift_float_conv3d(%arg0: tensor<1x3x4x3x3xf32>) -> (tensor<1x3x2x3x2xf32>) {
  %cst = "tf.Const"() {device = "", value = dense<1.0> : tensor<2x3x3x3x2xf32>} : () -> tensor<2x3x3x3x2xf32>
  %0 = "tf.Conv3D"(%arg0, %cst) {
    data_format = "NDHWC", device = "", dilations = [1, 1, 1, 1, 1], padding = "SAME", strides = [1, 1, 2, 1, 1]
  } : (tensor<1x3x4x3x3xf32>, tensor<2x3x3x3x2xf32>) -> tensor<1x3x2x3x2xf32>

  // Input: [N, H, W, C] for Conv2D or [N, D, H, W, C] for Conv3D.
  dnums.set_input_batch_dimension(0);
  dnums.set_input_feature_dimension(num_dims - 1);
  // Kernel: [K, K, I, O] for Conv2D or [K, K, K, I, O] for Conv3D.
  dnums.set_kernel_input_feature_dimension(num_dims - 2);
  dnums.set_kernel_output_feature_dimension(num_dims - 1);
  // Output: [N, H, W, C] for Conv2D or [N, D, H, W, C] for Conv3D.

   (IsInt8ElementType $filter),
   (IsConstTensor $filter),
   (IsInt32ElementType $conv),
   (HasStaticShapeConstraint $filter),
   (HasStaticShapeAtDimsConstraint<"3"> $input)],
  [], (addBenefit 10)>;

// Convert Conv2D with hybrid inputs (f32 activation/int8 weight) to XlaConv
def ConvertTFConv2DToXLAConvOpWeightOnly : Pat<
  (TF_Conv2DOp:$conv
    $input,
    (TF_MulOp (TF_CastOp (TF_IdentityOp $filter), $truncate1), $scale),

table Conv2DOptions {
  padding:Padding;
  stride_w:int;
  stride_h:int;
  fused_activation_function:ActivationFunctionType;
  dilation_w_factor:int = 1;
  dilation_h_factor:int = 1;
}

// Options for both Conv3D and Conv3DTranspose.
table Conv3DOptions {
  padding:Padding;
  stride_d:int;
  stride_w:int;
  stride_h:int;
  fused_activation_function:ActivationFunctionType;
  dilation_d_factor:int = 1;

  fused_activation_function:ActivationFunctionType;
  dilation_w_factor:int = 1;
  dilation_h_factor:int = 1;
  // Parameters for Conv2D version 8 or above.
  // When set, quantized_bias_type defines the dtype for both bias and accumulator.
  quantized_bias_type: TensorType;
}

// Options for both Conv3D and Conv3DTranspose.
table Conv3DOptions {
  padding:Padding;
  stride_d:int;
  stride_w:int;
  stride_h:int;