// 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-NEXT: return %0 : tensor<i1>
//CHECK-NEXT:}

func.func @dynamic_update_slice(%arg0: tensor<4x4xi64>, %arg1: tensor<2x3xi64>, %arg2: tensor<i64>, %arg3: tensor<i64>) -> tensor<4x4xi64> {

// CHECK:         tf.GlobalIterId
func.func @globaliterid_not_hoisted(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {
  %0:2 = "tf.WhileRegion"(%arg0, %arg1) ({
  ^bb0(%arg2: tensor<i32>, %arg3: tensor<i32>):
    %1 = "tf.OpA"() {is_stateless = true} : () -> tensor<i1>
    "tf.Yield"(%1) : (tensor<i1>) -> ()
  }, {
  ^bb0(%arg2: tensor<i32>, %arg3: tensor<i32>):
    %1 = "tf.GlobalIterId"() : () -> tensor<i64>

  // 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

func.func @add(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  %0 = "tf.Add"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  func.return %0: tensor<1xf32>

// CHECK-LABEL: add
// CHECK:  tfl.add %arg0, %arg1 {fused_activation_function = "NONE"} : tensor<1xf32>
// CHECK:  return
}

func.func @sub(%arg0: tensor<1xi64>, %arg1: tensor<1xi64>) -> tensor<1xi64> {

  %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: [[PI_0:%.*]] = "tf.TPUPartitionedInputV2"([[ARG0]], [[ARG1]])
  %pi_0 = "tf.TPUPartitionedInputV2"(%arg0, %arg1) {device = "", partition_dims = []} : (tensor<!tf_type.resource<tensor<10x3xf32>>>, tensor<!tf_type.resource<tensor<10x3xf32>>>) -> tensor<!tf_type.resource<tensor<10x3xf32>>>
  // CHECK: [[PI_1:%.*]] = "tf.TPUPartitionedInputV2"([[ARG2]], [[ARG3]])

// 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