int64_t channel = input_type.getDimSize(3);

  auto avg_pool_output_type = RankedTensorType::get(
      {batch, 1, 1, channel}, input_type.getElementType());
  auto avg_pool = rewriter.create<TFL::AveragePool2DOp>(
      mean_op.getLoc(), avg_pool_output_type, input,
      rewriter.getI32IntegerAttr(height), rewriter.getI32IntegerAttr(width),
      rewriter.getStringAttr("VALID"), rewriter.getI32IntegerAttr(1),

//===----------------------------------------------------------------------===//
// AveragePool2DOp
//===----------------------------------------------------------------------===//

int64_t AveragePool2DOp::GetArithmeticCount(Operation* op) {
  int64_t count;
  if (ArithmeticCountUtilHelper::GetFirstOutputCount(op, &count)) {
    auto avg_pool = llvm::dyn_cast<AveragePool2DOp>(op);