// CHECK:   call @while_body
// CHECK:   yield
// CHECK: while_body
// CHECK: while_cond

// INLINE: tfl.while
// INLINE: ^bb0([[ARGS:.*]]):
// INLINE:   tf.Less
// INLINE:   yield
// INLINE: ^bb0([[ARGS]]):
// INLINE:   %cst_2 = arith.constant
// INLINE:   yield
// INLINE-NOT: while_body
// INLINE-NOT: while_cond

// CANON-LABEL:   func @while_main
// CANON-SAME:         ([[VAL_0:%.*]]: tensor<?x256x256xf32>)

      "tf_device.cluster"() ( {
         %1 = "tf.While"(%0) {body = @while_body, cond = @while_cond}
         tf_device.return
      })
      return
    }
    func @while_body(%arg0: tensor<*x!tf_type.resource<tensor<f32>>>) {
     %constant = "tf.Const"() ...
     "tf.AssignVariableOp"(%arg0, %constant)
     return %arg0
    }
    func @while_cond(%arg0: tensor<*x!tf_type.resource<tensor<f32>>>) {
      %read = "tf.ReadVariableOp"(%arg0)

// RUN: tf-tfrt-opt -split-input-file -tf-mlrt-while-to-map-fn %s | FileCheck %s

// Test a while to map_fn conversion in which the max iteration is hard coded inside the predicate body.

// CHECK-LABEL: map/while_cond
func.func private @"map/while_cond"(%arg0: tensor<i32>, %arg1: tensor<i32>, %arg2: tensor<!tf_type.variant<tensor<*xf32>>>, %arg3: tensor<?xf32>) -> tensor<i1> {

    recompute_analysis_for_funcs.insert(
        while_op->getParentOfType<func::FuncOp>());
    func::FuncOp while_cond = while_op.cond_function();
    // Rewrite while op with extra chaining arguments and results.
    while_op = RewriteWhileOp(while_op, num_chains, chaining_data_type);
    bool first_visit = visited.insert(while_cond).second;
    if (!first_visit) continue;
    // Modify while condition function with extra chaining arguments.

// CHECK:             tfl.yield
// CHECK-SAME: (tensor<i32>, tensor<i32>, tensor<*xf32>, tensor<4x2xf32>, tensor<4x2xf32>, tensor<*xf32>, tensor<4x4x3xf32>) -> ()

// CHECK-LABEL:   func private @tfl.while_cond(