}

// CHECK-LABEL: reduce_window_not_4d
// CHECK: stablehlo.reduce_window
// CHECK-NOT: tfl.max_pool_2d

// -----

// Tests that a quantized `stablehlo.reduce_window` with max that takes multiple
// inputs is not converted to `tfl.max_pool_2d`.

func.func @reduce_window_not_binary(
  %arg0: tensor<3x2x9x10x3x!quant.uniform<i8:f32, 3.000000e-01:-5>>,

func.func @testMaxPool2D(tensor<256x32x32x3xf32>) -> tensor<?xf32> {
^bb0(%arg0: tensor<256x32x32x3xf32>):
  // CHECK: "tfl.max_pool_2d"(%arg0) <{filter_height = 1 : i32, filter_width = 1 : i32, fused_activation_function = "RELU6", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<256x32x32x3xf32>) -> tensor<?xf32>

      // iteration.
      input_dims.push_back(static_cast<int64_t>(dim_to_expand));
    }
    return input;
  }
};

// Rewrites quantized stablehlo.reduce_window with max to tfl.max_pool_2d.
class RewriteQuantizedReduceWindowOpWithMax
    : public OpRewritePattern<stablehlo::ReduceWindowOp> {
 public:
  using OpRewritePattern<stablehlo::ReduceWindowOp>::OpRewritePattern;

  %5 = arith.addf %0, %1 : tensor<1x1x1x16xf32>
  %6 = arith.addf %2, %5 : tensor<1x1x1x16xf32>
  func.return %6 : tensor<1x1x1x16xf32>

// CHECK-LABEL: func @maxPool2D
// CHECK:  "tfl.max_pool_2d"(%arg0) <{filter_height = 3 : i32, filter_width = 6 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 3 : i32, stride_w = 1 : i32}> : (tensor<1x1x1x16xf32>) -> tensor<1x1x1x16xf32>

    return quant::GetFixedOutputRange(is_signed, bit_width, result_type,
        /*scale=*/16.0 / 256, /*zero_point=*/127);
  }
  }];
}

def TFL_MaxPool2DOp : TFL_Op<"max_pool_2d", [
    TFL_OperandHasRank<0, 4>,
    PredOpTrait<"input and output must have same element type",
      TFL_TCresVTEtIsSameAsOp<0, 0>>,
    Pure,
    SameOperandsAndResultsScale,
    QuantizableResult,