}
  }
  attr {
    key: "body"
    value {
      func {
        name: "while_body_6"
      }
    }
  }
  attr {
    key: "cond"
    value {
      func {
        name: "while_cond_5"
      }
    }
  }
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
        }
        shape {
        }
        shape {
        }
      }

// CHECK-SAME: body = @while_body
// CHECK-SAME: cond = @while_cond
// CHECK: "tf.While"(%arg1)
// CHECK-SAME: body = @while_body_0
// CHECK-SAME: cond = @while_cond_1

// CHECK: func @while_body(%arg0: tensor<256x256xi32>)
// CHECK: func @while_cond(%arg0: tensor<256x256xi32>)
// CHECK: func private @while_body_0(%arg0: tensor<128xi32>)
// CHECK: func private @while_cond_1(%arg0: tensor<128xi32>)

    %0 = "tf.While"(%cst_main) {body = @while_body, cond = @while_cond, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
    // expected-error @+1 {{'tf.While' op multiple users of function.}}
    %1 = "tf.While"(%cst_main) {body = @while_body, cond = @while_cond, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
    return
  }

    %0 = "tf.While"(%cst_main) {body = @while_body, cond = @while_cond, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
    return
  }
  func.func private @while_body(%arg0: tensor<i32>) -> (tensor<i32>) {
    // Verify the overall pipelining control flow and supporting functions.
    // The order of these functions is also significant.
    // CHECK: {{.*StatefulPartitionedCall.* f = @while_cond.*}}

  }

  func.func private @while_cond(%arg0: tensor<i32> {tf._user_specified_name = "while/loop_counter"}, %arg1: tensor<i32> {tf._user_specified_name = "while/maximum_iterations"}, %arg2: tensor<1x5x5x1024xf32>) -> tensor<i1>
  attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<1x5x5x1024>], tf._original_func_name = "while_cond_60"} {

  // CHECK-NEXT: return %[[CONST]]
  func.return %arg0, %arg1 : tensor<*x!tf_type.resource<tensor<f32>>>, tensor<*x!tf_type.resource<tensor<f32>>>
}
// CHECK: func @while_cond1(%arg0: tensor<f32>)
func.func @while_cond1(%arg0: tensor<*x!tf_type.resource<tensor<f32>>>, %arg1: tensor<*x!tf_type.resource<tensor<f32>>>) -> tensor<f32> {

// CHECK-LABEL: func.func private @"map/while_cond"
func.func private @"map/while_cond"(%loop_count: tensor<i32>, %max_iterations: tensor<i32>, %handle: tensor<?x!tf_type.resource>, %flow_in: tensor<*xf32>, %matrix: tensor<3x3xf32>) -> tensor<i1> {
  %0 = "tf.Less"(%loop_count, %max_iterations) : (tensor<i32>, tensor<i32>) -> tensor<i1>
  return %0 : tensor<i1>
}

}

func.func @else(%x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> tensor<i1> {
  return %y: tensor<i1>
}

// CHECK-LABEL: func @while_cond_if
// CHECK: [[cond:%.*]] = tf_mlrt.predicate
// CHECK: [[z:%.*]] = mlrt.cond [[cond]] @then @else
// CHECK: return [[z]]
func.func @while_cond_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> (tensor<i1>) {

  // CHECK: return [[Result]]#0
  func.return %0#0 : tensor<*xf32>
}

// -----

// Simple trivially transformable while
// CHECK: func private @while_cond
// CHECK: func private @while_body
// CHECK-LABEL: testWhileRegionTrivial
func.func private @while_cond(%arg0 : tensor<*xf32>, %arg1 : tensor<i32>) -> tensor<i1>
func.func private @while_body(%arg0 : tensor<*xf32>, %arg1 : tensor<i32>) -> (tensor<*xf32>, tensor<i32>)

// RUN: tf-opt %s -tf-drop-while-shape-invariant | FileCheck %s
// RUN: tf-opt %s -tf-drop-while-shape-invariant-in-device-cluster | FileCheck -check-prefix=IN-CLUSTER %s


func.func @while_cond(%arg0: tensor<*xf32>) -> tensor<i1> {
  %0 = "tf.Const"() {value = dense<true> : tensor<i1>} : () -> tensor<i1>
  func.return %0 : tensor<i1>
}

func.func @while_body(%arg0: tensor<*xf32>) -> (tensor<*xf32>) {