%0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[1, 2, 3]> : tensor<3xi64>, name = "broadcast.0"}> : (tensor<8x1x16xf32>) -> tensor<3x8x8x16xf32>
  func.return %0 : tensor<3x8x8x16xf32>
}

// CHECK-LABEL:   func @broadcast_in_dim_general_case(
// CHECK-SAME:                                        %[[VAL_0:.*]]: tensor<3x1x16xf32>) -> tensor<3x8x8x16xf32> {

// Confirm a wide array of attribute survives the round-trip
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>

// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s

func.func @main(tensor<1x6x6x16xf32>) -> tensor<1x1x1x16xf32> {
^bb0(%arg0: tensor<1x6x6x16xf32>):
  // 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:   } ],

  func.return %1 : tensor<1x1x1x16xf32>

// CHECK: %[[avgp:.*]] = "tfl.average_pool_2d"(%arg0)
// CHECK: %[[dq:.*]] = "tfl.dequantize"(%[[avgp]]) : (tensor<1x1x1x16x!quant.uniform<u8:f32, 7.812500e-03:128>>) -> tensor<1x1x1x16xf32>
// CHECK: return %[[dq]] : tensor<1x1x1x16xf32>
}

// CHECK-LABEL: QuantizeReshape2D

  %0 = "tfl.dequantize"(%arg0) : (tensor<1x6x6x16x!quant.uniform<u8:f32, 0.1>>) -> tensor<1x6x6x16xf32>
  %1 = "tfl.dequantize"(%arg1) : (tensor<1x6x6x16x!quant.uniform<u8:f32, 0.1>>) -> tensor<1x6x6x16xf32>
  %2 = "tfl.minimum"(%0, %1) : (tensor<1x6x6x16xf32>, tensor<1x6x6x16xf32>) -> tensor<1x6x6x16xf32>
  func.return %2 : tensor<1x6x6x16xf32>

// CHECK: %0 = "tfl.dequantize"(%arg0)
// CHECK: %1 = "tfl.dequantize"(%arg1)

^bb0(%arg0: tensor<1x7x7x16xf32>):
  // expected-error @+1 {{requires attribute 'padding'}}
  %0 = "tf.AvgPool"(%arg0) {T = "tfdtype$DT_FLOAT", ksize = [1, 7, 7, 1], strides = [1, 1, 1, 1]} : (tensor<1x7x7x16xf32>) -> tensor<1x1x1x16xf32>
  func.return %0 : tensor<1x1x1x16xf32>
}

// -----

func.func @testAvgPoolWrongPadding(tensor<1x7x7x16xf32>) -> tensor<1x1x1x16xf32> {
^bb0(%arg0: tensor<1x7x7x16xf32>):

func.func @partitioned_3d(%arg0: tensor<!tf_type.resource<tensor<16x8x16xf32>>>, %arg1: tensor<!tf_type.resource<tensor<16x8x16xf32>>>, %arg2: tensor<!tf_type.resource<tensor<16x8x16xf32>>>, %arg3: tensor<!tf_type.resource<tensor<16x8x16xf32>>>)...

  // Unsupported strides
  %2 = "tf.AvgPool"(%arg0) {T = "tfdtype$DT_FLOAT", data_format = "NHWC", ksize = [1, 3, 6, 1], padding = "VALID", strides = [1, 3, 1, 3]} : (tensor<1x6x6x16xf32>) -> tensor<1x1x1x16xf32>

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

  func.return %0 : tensor<16x8x8x1xf32>
  }
}

// CHECK: module {
// CHECK-NEXT:   func.func @main(%arg0: tensor<8x8x1x207xf32>, %arg1: tensor<3x3x16x207xf32>) -> tensor<16x8x8x1xf32> {

  %0:4 = "tfl.split"(%arg0, %arg1) {num_splits = 4 : i32, tac.device = "CPU"} : (tensor<i32>, tensor<1x8x8x1024xf32>) -> (tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>)
  func.return %0#0, %0#1, %0#3 : tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>
}