// tfl.average_pool_2d, this function finds the return type of the needed
// tfl.pad .
ShapedType GetPadOpType(mhlo::CompositeOp op);

// Given a Composite op that wraps a core.aten.avg_pool2d, finds the padding
// attribute to be passed to the a tfl.average_pool_2d that can fully replace
// this composite (here, padding is done directly by the tfl.average_pool_2d as

  func.return %0 : tensor<1x1x1x512xf32>
}

// CHECK:       func @mean_4d_keepdim([[VAL_0:%.*]]: tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32> {
// CHECK:           [[VAL_1:%.*]] = "tfl.average_pool_2d"([[VAL_0]]) <{filter_height = 48 : i32, filter_width = 48 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32>

func.func @main(tensor<1x6x6x16xf32>) -> tensor<1x1x1x16xf32> {
^bb0(%arg0: tensor<1x6x6x16xf32>):
  // CHECK: "tfl.average_pool_2d"(%{{.*}}) <{filter_height = 3 : i32, filter_width = 6 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 3 : i32, stride_w = 1 : i32}> : (tensor<1x6x6x16xf32>) -> tensor<1x1x1x16xf32>

// Replaces an aten.avg_pool2d with a (T -> tfl.average_pool_2d  -> T).
// Constraints are added on the attributes of the aten.avg_pool2d to ensure only
// ops that match the behaviour of tfl.average_pool_2d are directly lowered.
def LegalizeAvgPool2dComposite: Pat<
                    (MHLO_CompositeOp:$old_val
                    (variadic $a_input), 

// Returns a NHWC shaped type from an NCHW shaped type op.
// For example- Given a Composite op that wraps a core.aten.avg_pool2d, this
// returns the return type of the tfl.average_pool_2d emitted. Note that the
// aten.avg_pool2d works with the NCHW layout while tfl.average_pool_2d assumes
// NHWC.
ShapedType GetNhwcReturnTypeFromNchw(Operation* old_op);

}  // namespace odml

}  // namespace mlir

  // CHECK:      {
  // CHECK-NEXT:   version: 3,
  // CHECK-NEXT:   operator_codes: [ {
  // CHECK-NEXT:     deprecated_builtin_code: 1,
  // CHECK-NEXT:     version: 1,
  // CHECK-NEXT:     builtin_code: AVERAGE_POOL_2D
  // CHECK-NEXT:   } ],
  // CHECK-NEXT:   subgraphs: [ {
  // CHECK-NEXT:     tensors: [ {
  // CHECK-NEXT:       shape: [ 1, 6, 6, 16 ],
  // CHECK-NEXT:       buffer: 1,
  // CHECK-NEXT:       name: "arg0",

  func.return %7 : tensor<1x5x5x3xf32>
// DEBUG: %[[q:.*]] = "tfl.quantize"(%arg0)
// DEBUG: %[[dq:.*]] = "tfl.dequantize"(%[[q]])
// DEBUG: "tfl.average_pool_2d"(%[[dq]])
// DEBUG: "tfl.average_pool_2d"(%[[q]])
}

// MODEL-DEBUG-LABEL: CheckNumericVerifyWholeModel
func.func @CheckNumericVerifyWholeModel(%arg0: tensor<1x4x4x3xf32>) -> tensor<1x1x1x3xf32> {

// CHECK:           %[[VAL_4:.*]] = "tfl.pad"(%[[VAL_2]], %[[VAL_3]]) : (tensor<1x6x6x3xf32>, tensor<4x2xi32>) -> tensor<1x6x6x3xf32>
// CHECK:           %[[VAL_5:.*]] = "tfl.average_pool_2d"(%[[VAL_4]]) <{filter_height = 3 : i32, filter_width = 3 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x6x6x3xf32>) -> tensor<1x4x4x3xf32>

}

// -----

func.func @testAvgPool(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
  // CHECK: tac.device = "GPU", tac.inference_type = "FLOAT"
  %0 = "tfl.average_pool_2d"(%arg0) {filter_height = 3 : i32, filter_width = 3 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32>

}

func.func @testAveragePool2D(tensor<1x10x10x3xf32>) -> tensor<1x10x10x3xf32> {
^bb0(%arg0: tensor<1x10x10x3xf32>):
  // CHECK: _arithmetic_count = 2700 : i64
  %0 = "tfl.average_pool_2d"(%arg0) {filter_height = 3 : i32, filter_width = 3 : i32, fused_activation_function = "RELU6", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32} : (tensor<1x10x10x3xf32>) -> tensor<1x10x10x3xf32>