%out_scale : tensor<*xf32>, %out_zp : tensor<*xi32>) -> tensor<*x${output_type}>
    attributes {tf_quant.quantized_ops = ${quantized_ops}} {
      // Given the convolution takes 2 qint8 inputs and output a qint32.
      // The accumulation scale is (input_scale * filter_scale).
      // The accumulation zero point is 0.

    %0 = stablehlo.convolution(%arg0, %arg1) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = [[0, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
    return %0 : tensor<1x3x4x2xf32>
  }
  // CHECK: func private @composite_conv_fn
  // CHECK: %[[CONV:.+]] = stablehlo.convolution
  // CHECK: return %[[CONV]]
}

// usually has the following pattern. In the example below,
// the input operand would be stablehlo.convolution op, and return value would
// be stablehlo.add op.
//
// ```
// %0 = stablehlo.constant dense<3>
// %1 = stablehlo.constant dense<4>
// %2 = stablehlo.constant dense<2>
// %3 = stablehlo.convolution(%%arg0, %%arg1) :
//          (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<?x3x4x2xf32>

  // CHECK: %[[CONST:.*]] = mhlo.constant()
  // CHECK-SAME{LITERAL} value = dense<127> : tensor<2x3x3x2xi8>
  // CHECK-SAME: tensor<2x3x3x2x!quant.uniform<i8:f32, 1.000000e+00:3>>
  // CHECK: mhlo.convolution(%arg0, %[[CONST]])
  // CHECK-SAME{LITERAL}: dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
  // CHECK-SAME{LITERAL}: window = {stride = [1, 2], pad = [[0, 0], [0, 0]], lhs_dilate = [1, 1], rhs_dilate = [2, 2]}

class DeferActivationTransposeForAddOp : public OpRewritePattern<AddOp> {
 public:
  using OpRewritePattern<AddOp>::OpRewritePattern;

  LogicalResult match(AddOp op) const override {
    // Only supports the case for 2D convolution.
    const Value lhs = op.getOperand(0);
    if (!HasRankOf(lhs, /*rank=*/4)) return failure();

    const Value rhs = op.getOperand(1);
    Operation* rhs_op = rhs.getDefiningOp();

        ).astype('f4')

      @def_function.function
      def conv2d(self, input_tensor: core.Tensor) -> Mapping[str, core.Tensor]:
        """Performs a 2D convolution operation.

        Args:
          input_tensor: Input tensor to perform convolution on.

        Returns:
          A map of: output key -> output result.
        """
        scale = [1.0] * self.out_channel_size

  EXPECT_THAT(default_spec.matcher().function_name().regex(), StrEq(".*"));
  EXPECT_TRUE(default_spec.method().has_static_range_ptq());

  // Test that the expansion for convolution ops is done.
  const QuantizationSpec& conv_spec = new_config.specs().specs(1);
  EXPECT_THAT(conv_spec.matcher().function_name().regex(),
              StrEq("composite_conv.*"));

inline constexpr std::array<int64_t, 4> kNchwToNhwcPermutation = {0, 2, 3, 1};

// Permutation from the OIHW (== (output features, input features, height,
// width)) tensor format to HWIO. This is commonly used to transpose convolution
// weights represented as OIHW format to HWIO, which is more desirable for
// certain downstream optimization passes (e.g. XLA).
inline constexpr std::array<int64_t, 4> kOihwToHwioPermutation = {2, 3, 1, 0};

    return true;
  if (op_name == "stablehlo.broadcast_in_dim" &&
      field_name == "broadcast_dimensions")
    return true;
  if ((op_name == "stablehlo.convolution" ||
       op_name == "stablehlo.dynamic_conv") &&
      (field_name == "window_strides" || field_name == "lhs_dilation" ||
       field_name == "rhs_dilation"))
    return true;

    value = builder.create<TF::ReshapeOp>(
        loc, value, Create1DConstValue(builder, loc, new_shape));
  }
  return ConstantFoldOpIfPossible(value.getDefiningOp()).front();
}

// Matches convolution op with "NHWC" data format or matmul op with false adj_y.
// The list of supported ops in this function is:
// - Conv2DOp
// - Conv3DOp
// - DepthwiseConv2dNativeOp
// - MatMulOp
// - BatchMatMulV2Op