int64_t MaxPool2DOp::GetArithmeticCount(Operation* op) {
  int64_t count;
  if (ArithmeticCountUtilHelper::GetFirstOutputCount(op, &count)) {
    auto max_pool = llvm::dyn_cast<MaxPool2DOp>(op);
    return max_pool.getFilterHeight() * max_pool.getFilterWidth() * count;
  }

  return -1;
}

//===----------------------------------------------------------------------===//
// L2NormalizationOp

// CHECK: %[[MAX_POOL:.*]] = "tfl.max_pool_2d"(%[[ARG0]])
// CHECK-SAME: {filter_height = 3 : i32, filter_width = 4 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 2 : i32, stride_w = 3 : i32}
// CHECK-SAME: (tensor<2x9x10x3x!quant.uniform<i8:f32, 3.000000e-01:-5>>) -> tensor<2x4x3x3x!quant.uniform<i8:f32, 3.000000e-01:-5>>
// CHECK: return %[[MAX_POOL]]

// -----

// "SAME").
bool IsSpatialPoolingWithoutDilation(
    mhlo::ReduceWindowOp rw, llvm::SmallVectorImpl<int64_t>* window_strides,
    std::string* padding_mode, std::string* data_format) {
  // tf.max_pool or tf.avg_pool need at least 3 dimensions (batch, spatial,
  // channel).
  const uint64_t rank = rw.getWindowDimensions().size();
  if (rank <= 3 || rank > 5) return false;

// CHECK:  %1 = "tf.MaxPool"(%arg0)
// CHECK:  %2 = "tf.MaxPool"(%arg0)
}

func.func @maxPool2DChannelFirst(%arg0: tensor<1x16x6x6xf32>) -> tensor<1x16x1x1xf32> {
  // OK