}
    }
  }
  attr {
    key: "_lower_using_switch_merge"
    value {
      b: true
    }
  }
  attr {
    key: "body"
    value {
      func {
        name: "while_body_6"
      }
    }
  }
  attr {
    key: "cond"
    value {
      func {
        name: "while_cond_5"
      }
    }
  }
  attr {
    key: "output_shapes"
    value {

// CHECK: %0 = "tf.While"(%arg0)
// 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>)

    // CHECK: {{.*tf.While.*body = @while_body.* cond = @while_cond.* parallel_iterations = 3}}
    %0 = "tf.While"(%cst) {body = @while_body, cond = @while_cond, is_stateless = false, parallel_iterations = 3} : (tensor<i32>) -> (tensor<i32>)
    return
  }
  func.func private @while_body(%arg0: tensor<i32>) -> (tensor<i32>) {
    // metadata ops

    %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
  }
  func.func private @while_body(%arg0: tensor<i32>) -> (tensor<i32>) {

  func.return %1 : tensor<i1>
}

// The @while_body function is a Multi-hosts function which contains three
// parts:
//   - A local subgraph which contains both local ops and remote_run kernels to
//     trigger remote subgraphs
//   - Two remote subgraph which contains remote ops on worker:1 and worker:2.
// CHECK: func @while_body(%[[ARG_0:.*]]: tensor<i32> {tf.device = "/job:localhost/replica:0/task:0/device:CPU:0"})

  func.return
}
// CHECK: func @while_body(%[[BARG0:.*]]: tensor<f32>)
func.func @while_body(%arg0: tensor<*x!tf_type.resource<tensor<f32>>>, %arg1: tensor<*x!tf_type.resource<tensor<f32>>>)
    -> (tensor<*x!tf_type.resource<tensor<f32>>>, tensor<*x!tf_type.resource<tensor<f32>>>) {
  // CHECK: %[[WHILE:.*]] = "tf.While"(%[[BARG0]])
  %0:2 = "tf.While"(%arg0, %arg1) {

    %cst = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %cst_0 = "tf.Const"() {value = dense<1.0> : tensor<1x5x5x1024xf32>} : () -> tensor<1x5x5x1024xf32>
    // Check that these constants are remained in constants.
    // CHECK: func private @while_body

// CHECK-NEXT:    %0 = "tf.Less"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i1>
// CHECK-NEXT:    return %0 : tensor<i1>

// CHECK-LABEL: func.func private @"map/while_body"

  return %r : tensor<i1>
}

// CHECK-LABEL: func @while_body_if
func.func @while_body_if(%cond: tensor<i1>, %x: tensor<i1>, %y: tensor<i1>, %z: tensor<i32>) -> (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i32>) {

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>) {
  %0 = "tf.SomeOp"(%arg0) : (tensor<*xf32>) -> tensor<*xf32>
  func.return %0 : tensor<*xf32>
}

// Test that -tf-drop-while-shape-invariant-in-device-cluster pass does not drop