// CHECK:           %[[VAL_13:.*]] = "tfl.conv_2d"(%[[VAL_11]], %[[VAL_12]], %[[VAL_4]]) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x384x512xf32>, tensor<128x1x1x512xf32>, tensor<128xf32>) -> tensor<1x1x384x128xf32>

func.func @mean_4d_keepdim(%arg0: tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32> {
  %cst = arith.constant dense<[1, 2]> : tensor<2xi32>
  %0 = "tfl.mean"(%arg0, %cst) {keep_dims = true} : (tensor<1x48x48x512xf32>, tensor<2xi32>) -> tensor<1x1x1x512xf32>
  func.return %0 : tensor<1x1x1x512xf32>
}

// CHECK:       func @mean_4d_keepdim([[VAL_0:%.*]]: tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32> {

  %w_1 = arith.constant dense<127.0> : tensor<4096x1x1x1xf32>
  %w_2 = arith.constant dense<127.0> : tensor<128x1x1x1xf32>
  %b = arith.constant dense<127.0> : tensor<2048x1x1x1xf32>
  %custom_1 = "tfl.custom"(%0, %w_1, %w_2, %b) {custom_code = "CustomTestOp", custom_option = #tfl<const_bytes : "0x">} : (tensor<1x1x1x1xf32>, tensor<4096x1x1x1xf32>, tensor<128x1x1x1xf32>, tensor<2048x1x1x1xf32>) -> tensor<*xf32>

  %w_1 = arith.constant dense<127.0> : tensor<4096x1x1x1xf32>
  %w_2 = arith.constant dense<127.0> : tensor<128x1x1x1xf32>
  %b = arith.constant dense<127.0> : tensor<2048x1x1x1xf32>
  %custom_1 = "tfl.custom"(%0, %w_1, %w_2, %b) {custom_code = "CustomTestOp", custom_option = #tfl<const_bytes : "0x">} : (tensor<1x1x1x1xf32>, tensor<4096x1x1x1xf32>, tensor<128x1x1x1xf32>, tensor<2048x1x1x1xf32>) -> tensor<*xf32>