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

          (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

    // than function name.
    if (!function_name.starts_with("composite_")) {
      return failure();
    }

    if (function_name.contains("depthwise_conv2d")) {
      // Uniform Quantized op requires weights of tf.DepthwiseConv2dNative to
      // be transformed from [H,W,C,M] to [H,W,1,CxM] where
      // H=height,W=width,C=channel,M=multiplier. Therefore, a reshape op is