return broadcast_in_dim_op.getOperand().getDefiningOp();
  }
  return op;
}

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

  auto func_op = module_op->lookupSymbol<func::FuncOp>("fully_quantized_add");
  ASSERT_THAT(func_op, NotNull());

  auto add_op_itr = func_op.getBody().op_begin<mlir::stablehlo::AddOp>();
  ASSERT_THAT(add_op_itr,
              Ne(func_op.getBody().op_end<mlir::stablehlo::AddOp>()));

  EXPECT_TRUE(IsOpFullyQuantized(*add_op_itr));
}

TEST_F(IsOpFullyQuantizedTest, FalseIfOpNotQuantized) {

  // explicit activation clipping op is not required.
  if (isa<AddOp>(next_op) && gemm_style_op->hasOneUse()) {
    // bias fusion
    CreateAndReturnQuantizedBiasPattern(
        next_op, rewriter, entry_func_op, func_result_type,
        accumulation_quantized_element_type, gemm_style_op);
  } else if (auto add_op = cast_or_null<AddOp>(
                 GetBroadcastedUserOp<AddOp>(gemm_style_op))) {
    // broadcasted bias fusion

// Fuse Add with proceeding FullyConnected.
// TODO(b/136285429): Move to tablegen when variadic is supported
struct FuseFullyConnectedAndAdd : public OpRewritePattern<TFL::AddOp> {
  using OpRewritePattern<TFL::AddOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(TFL::AddOp add_op,
                                PatternRewriter &rewriter) const override {
    // Match Add.
    DenseElementsAttr added_value;

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tsl/platform/status_matchers.h"

namespace mlir::quant {
namespace {

using ::mlir::stablehlo::AddOp;
using ::mlir::stablehlo::ConstantOp;
using ::mlir::stablehlo::ConvolutionOp;
using ::mlir::stablehlo::DotGeneralOp;
using ::mlir::stablehlo::SubtractOp;
using ::testing::ElementsAreArray;
using ::testing::Eq;

 public:
  using OpRewritePattern<stablehlo::AddOp>::OpRewritePattern;

  LogicalResult match(stablehlo::AddOp op) const override {
    return success(IsI8F32UniformQuantizedType(GetElementType(op.getLhs())) &&
                   IsI8F32UniformQuantizedType(GetElementType(op.getRhs())));
  }

  void rewrite(stablehlo::AddOp op, PatternRewriter& rewriter) const override {

                                                  one_over_root2_cst);
  auto erf_op = rewriter.create<chlo::ErfOp>(mul_op.getLoc(), mul_op);
  auto add_op =
      rewriter.create<stablehlo::AddOp>(erf_op.getLoc(), erf_op, one_cst);
  auto lhs_mul_op = rewriter.create<stablehlo::MulOp>(
      half_cst.getLoc(), new_func.getArguments()[0], half_cst);

  Value filter_size = CreateScalarConstValue<int32_t>(builder, loc, filter_sz);
  Type int32_scalar_type = zero.getType();

  auto scalar_add = [&](Value lhs, Value rhs) {
    return builder.create<TF::AddOp>(loc, int32_scalar_type, lhs, rhs);
  };
  auto scalar_mul = [&](Value lhs, Value rhs) {
    return builder.create<TF::MulOp>(loc, int32_scalar_type, lhs, rhs);
  };
  auto scalar_sub = [&](Value lhs, Value rhs) {

        bn_op.getVariance(), variance_type, rewriter);
    if (!epsilon) {
      return failure();
    }

    // Compute multiplier = scale / sqrt(variance + epsilon)
    Value multiplier = rewriter.create<mhlo::AddOp>(
        bn_op.getLoc(), bn_op.getVariance(), epsilon);
    multiplier = rewriter.create<mhlo::RsqrtOp>(bn_op.getLoc(), multiplier);
    multiplier = rewriter.create<mhlo::MulOp>(bn_op.getLoc(), multiplier,

  return std::make_unique<CpuConcatOp>();
}

TAC_REGISTER_CPU_OP(Conv2DOp, CreateConvOp);
TAC_REGISTER_CPU_OP(DepthwiseConv2DOp, CreateConvOp);
TAC_REGISTER_CPU_OP(FullyConnectedOp, CreateConvOp);
TAC_REGISTER_CPU_OP(AddOp, CreateArithmeticOp);
TAC_REGISTER_CPU_OP(MulOp, CreateArithmeticOp);
TAC_REGISTER_CPU_OP(ConcatenationOp, CreateConcatOp);
TAC_REGISTER_CPU_OP(ReshapeOp, CreateConcatOp);
TAC_REGISTER_CPU_OP(PackOp, CreateConcatOp);