// CHECK-LABEL: func @FusedBatchNormRelu
func.func @FusedBatchNormRelu(%arg0: tensor<8x8x8x8xf32>, %arg1: tensor<8xf32>, %arg2: tensor<8xf32>, %arg3: tensor<8xf32>, %arg4: tensor<8xf32>) -> (tensor<8x8x8x8xf32>) {
// CHECK-NEXT: %[[Y:[a-z0-9]*]], {{.*}}_FusedBatchNormEx
// CHECK-NEXT: return %[[Y]]

// CHECK-LABEL: func @__forward_call_369
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain
// CHECK-SAME: [[arg1:%.*]]: !tfrt_fallback.tf_tensor {tf._user_specified_name = "inputs"},
// CHECK-SAME: [[arg2:%.*]]: !tfrt_fallback.tf_tensor, [[arg3:%.*]]: !tfrt_fallback.tf_tensor, [[arg4:%.*]]: !tfrt_fallback.tf_tensor, [[arg5:%.*]]: !tfrt_fallback.tf_tensor)
// CHECK-SAME: -> (!tfrt.chain
// CHECK: [[o1:%.*]] = tfrt_fallback_async.const_dense_tensor

// CHECK-SAME: (%[[ARG0:.*]]: tensor<f32>, %[[ARG1:.*]]: tensor<f32>, %[[ARG2:.*]]: tensor<*xf32>)
func.func @test_supported_lowering_of_tf_while_region(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<*xf32>) -> (tensor<f32>){
  %0:2 = "tf.WhileRegion"(%arg0, %arg2) ({
  ^bb0(%arg3: tensor<f32>, %arg4: tensor<*xf32>):
    %1 = "tf.Identity"(%arg3) : (tensor<f32>) -> tensor<f32>

// CHECK: func @while_loop_test(%[[ARG_0:.*]]: tensor<i32>, %[[ARG_1:.*]]: tensor<i32>, %arg2: tensor<*xf32>)
func.func @while_loop_test(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<*xf32>) {
  // CHECK-NEXT: tf.WhileRegion
  %0:2 = "tf.WhileRegion"(%arg0, %arg2) ({
  // CHECK-NEXT: bb0(%[[ARG_3:.*]]: tensor<i32>, %[[ARG_4:.*]]: tensor<*xf32>)
  ^bb0(%arg3: tensor<i32>, %arg4: tensor<*xf32>):
    // CHECK-NEXT: %[[RESULT_1:.*]] = "tf.Identity"(%[[ARG_3]])

  // CHECK: func @f
  // CHECK: %arg0: tensor<!tf_type.resource<tensor<f32>>> {tf_saved_model.bound_input = @v}
  // CHECK: %arg1: tensor<!tf_type.resource<tensor<f32>>> {tf_saved_model.bound_input = @w}
  // CHECK: %arg2: tensor<!tf_type.resource<tensor<f32>>> {tf_saved_model.bound_input = @x}
  // CHECK-NOT: %arg3
  // CHECK-NOT: %arg4
  func.func @f(

  // CHECK: "tfl.conv_2d"
}

// CHECK-LABEL: fused_batch_norm_v3_bf16
func.func @fused_batch_norm_v3_bf16(%arg0: tensor<8x8x8x8xbf16>, %arg1: tensor<8xf32>, %arg2: tensor<8xf32>, %arg3: tensor<8xf32>, %arg4: tensor<8xf32>) -> tensor<8x8x8x8xbf16> {

// Ensure basic_lstm roundtrip exactly

func.func @main(%arg0: tensor<1x384xf32>, %arg1: tensor<1x96xf32>, %arg2: tensor<384x480xf32>, %arg3: tensor<384xf32>, %arg4: tensor<1x96xf32>) -> tensor<1x96xf32> {
// CHECK-LABEL: @main

  %7 = "tf.AddV2"(%arg4, %6) {T = f32, device = ""} : (tensor<*xf32>, tensor<?x?xf32>) -> tensor<*xf32>
  %8 = "tf.AddV2"(%arg0, %cst_2) {T = i32, device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
  func.return %8, %arg1, %0, %arg3, %7 : tensor<*xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*xf32>, tensor<*xf32>
}

// CHECK-LABEL: func @transpose_resnet_layer
func.func @transpose_resnet_layer(%arg0: tensor<?x224x224x3xf32>, // input
                             %arg1: tensor<64xf32>,          // batch_norm args
                             %arg2: tensor<256xf32>,          // batch_norm args
                             %arg3: tensor<7x7x3x64xf32>,    // conv filter #0
                             %arg4: tensor<1x1x64x256xf32>   // conv filter #1

  %cst = arith.constant dense<1> : tensor<i32> loc("dec")
  %arg0 = arith.constant dense<5> : tensor<i32> loc("N")
  %arg1 = arith.constant dense<3.0> : tensor<1xf32> loc("val")
  %0:3 = "tfl.while"(%arg0, %arg1, %cst) ({
    ^bb0(%arg2: tensor<*xi32>, %arg3: tensor<*xf32>, %arg4: tensor<i32>):
      %cst_0 = arith.constant dense<0> : tensor<i32>
      %1 = "tfl.greater"(%arg2, %cst_0) : (tensor<*xi32>, tensor<i32>) -> tensor<i1>