return %0 : tensor<1x1x2x2xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_4(%arg0: tensor<1x1x3x3xf32>) -> tensor<1x1x2x2xf32>

// CHECK-LABEL: avg_pool2d_5
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %0 = "tfl.transpose"(%arg0, %cst) : (tensor<1x1x3x3xf32>, tensor<4xi32>) -> tensor<1x3x3x1xf32>

// CHECK: return [[MUL1]]
}

func.func @batchmatmul2fullyconnected(%arg0: tensor<4x128x2xf32>) -> (tensor<4x128x1xf32>) {
  %0 = "tf.Const"() {value = dense<[[1.0], [2.0]]> : tensor<2x1xf32>} : () -> tensor<2x1xf32>
  %1 = "tf.BatchMatMulV2"(%arg0, %0) : (tensor<4x128x2xf32>, tensor<2x1xf32>) -> tensor<4x128x1xf32>
  func.return %1 : tensor<4x128x1xf32>

  // CHECK-LABEL: batchmatmul2fullyconnected

}

// -----

func.func @testPReluInvalidBroadcast(%arg0: tensor<15x14x2x14xf32>, %arg1: tensor<1x1x3xf32>) -> tensor<15x14x2x14xf32> {
  // expected-error @+1 {{'tfl.prelu' op operands don't have broadcast-compatible shapes}}
  %0 = "tfl.prelu"(%arg0, %arg1) : (tensor<15x14x2x14xf32>, tensor<1x1x3xf32>) -> tensor<15x14x2x14xf32>
  func.return %0 : tensor<15x14x2x14xf32>
}

// -----