attrs=[
        'stride_w: int', 'stride_h: int', 'filter_width: int',
        'filter_height: int', 'padding: {"SAME", "VALID"}'
    ],
    derived_attrs=['T: {float, int8}'],
    outputs=['o: T'])
def _composite_max_pool(input_, stride_w, stride_h, filter_width, filter_height,
                        padding):
  ksize = [1, filter_width, filter_height, 1]
  strides = [1, stride_w, stride_h, 1]

}

// 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>
// CHECK:           return [[VAL_1]] : tensor<1x1x1x512xf32>

^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>

  // CHECK-NEXT:       builtin_options: {
  // CHECK-NEXT:         padding: VALID,
  // CHECK-NEXT:         stride_w: 1,
  // CHECK-NEXT:         stride_h: 3,
  // CHECK-NEXT:         filter_width: 6,
  // CHECK-NEXT:         filter_height: 3
  // CHECK-NEXT:       }
  // CHECK-NEXT:     } ]
  // CHECK-NEXT:    name: "main"
  // CHECK-NEXT:   } ],
  // CHECK-NEXT:   description: "MLIR Converted.",

}

// -----

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.return %0 : tensor<256x30x30x16xf32>
}

// -----

}

func.func @testMaxPool2D(tensor<1x10x10x3xf32>) -> tensor<1x10x10x3xf32> {
^bb0(%arg0: tensor<1x10x10x3xf32>):
  // CHECK: _arithmetic_count = 2700 : i64
  %0 = "tfl.max_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>
  func.return %0 : tensor<1x10x10x3xf32>
}

                                    ConstantAttr<RankedI32ElementsAttr<[4]>,"{0, 2, 3, 1}">)),
                            /*filter_height*/(GetI32At<0> (GetAsVectorAttr<"kernel_size"> $attrs)),
                            /*filter_width*/(GetI32At<1> (GetAsVectorAttr<"kernel_size"> $attrs)),
                            /*padding*/(GetAvgPoolOpPadAttr $old_val),
                            /*stride_h*/(GetI32At<0> (GetAsVectorAttr<"stride"> $attrs)),

// 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>

  def test_new_max_pool(self):
    input_ = tf.random.uniform([8, 4, 4, 1])
    kwargs = {
        'input_': input_,
        'stride_w': 2,
        'stride_h': 2,
        'filter_width': 1,
        'filter_height': 1,
        'padding': 'SAME',
    }

    self._assertOpAndComposite([input_],
                               tf.function(gen_mnist_ops.new_max_pool),

              IsIntList1XY1:$strides,
              $padding,
              IsDataFormatNHWC:$format),
          (TFL_AveragePool2DOp $value,
              /*filter_height=*/ExtractI32At<1>:$ksize,
              /*filter_width=*/ExtractI32At<2>:$ksize,
              /*padding=*/$padding,
              /*stride_h=*/ExtractI32At<1>:$strides,
              /*stride_w=*/ExtractI32At<2>:$strides,