auto ret1 = ops::_Retval(s.WithOpName("ret1"), if_op.output[1], 1);
  NameAttrList cond_fn, body_fn;
  cond_fn.set_name("f3");
  body_fn.set_name("f2");
  auto while_op =
      ops::While(s.WithOpName("while"),
                 std::initializer_list<Input>{arg0, arg1}, cond_fn, body_fn);
  auto ret2 = ops::_Retval(s.WithOpName("ret2"), while_op.output[0], 2);

                                           Output* output) {
    *output = ToOutput(while_ctx->cond_output());
    return absl::OkStatus();
  };

  // Body function that adds one to input.
  BodyGraphBuilderFn body_fn = [](const Scope& scope,
                                  const std::vector<Output>& inputs,
                                  std::vector<Output>* outputs) {
    DCHECK_EQ(inputs.size(), 1);

    results: a list of futures.
    body_fn: input is [predicate_promise, arg0_future, arg0_promise, arg1_future, arg1_promise, ..., invariant_args] and has no returned results.
  }];

  let arguments = (ins TF_Tensor:$cond,
                       Variadic<TF_Tensor>:$arguments,
                       FlatSymbolRefAttr:$body_fn,
                       I32Attr:$invariant_size);

      last element of results is an I1 value that is false.
    body_fn: The body function that takes the arguments and returns the results
      that includes an I1 value to indicate whether next iteration should be executed.

    The pseudo code:

    while(cond) {
      results = body_fn(operands)
      cond = results#1
    }
    return results

  }];

    results: a list of futures.
    body_fn: its input are [predicate_promise, arg0_future, arg0_promise, arg1_future, arg1_promise, ..., invariant_args] and has no returned results.
  }];

  let arguments = (ins TFTensorType:$cond,
                       Variadic<TFTensorType>:$arguments,
                       FlatSymbolRefAttr:$body_fn,
                       I32Attr:$invariant_size);

    TF_CHECK_OK(GraphToFunctionDef(*g, "body_fn", body_fn_fdef));
  }
  {
    tensorflow::Scope s = tensorflow::Scope::NewRootScope();
    Output input = ops::Const(s.WithOpName("const0"), true, {2});
    NameAttrList cond_fn;
    cond_fn.set_name("cond_fn");
    NameAttrList body_fn;
    body_fn.set_name("body_fn");
    auto while_op =

static LogicalResult LowerWhileOp(WhileOp op) {
  Operation* op_inst = op.getOperation();
  Location loc = op_inst->getLoc();

  OpBuilder builder(op_inst);

  auto cond_fn = op.cond_function();
  auto body_fn = op.body_function();

  // Split the block containing the While op into two blocks.  One containing
  // operations before the While op and other containing the rest.  Create two

  %max_iter = "tf.Const"() {__op_key = 1, value = dense<2> : tensor<i32>} : () -> tensor<i32>
  // CHECK: tf_mlrt.map_fn
  %result = "tf_mlrt.tf_map_fn"(%max_iter, %input2, %2) { operandSegmentSizes = array<i32: 1, 1, 1>, body_fn = @NopMapFnBody, num_tensor_list_or_flow_in = 1 : i32} : (tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>) -> tensor<i32>
  return %result : tensor<i32>
}

// CHECK-LABEL: @NopMapFnBody

  // CHECK: [[map_fn_result:%.*]] = tf_mlrt.tf_map_fn([[reserve_size]], [[tensor_list]], %arg0)
  // CHECK-SAME: {body_fn = @"map/while_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
  // CHECK-NOT: tf.While

            /* nreturn_nodes */ 1, &cond_output));
        *output = cond_output[0];
        return absl::OkStatus();
      };

  // 'body_fn' copies the body graph into the parent graph.
  tensorflow::ops::BodyGraphBuilderFn body_fn =
      [params, parent, num_loop_vars](
          const tensorflow::Scope& scope,
          const std::vector<tensorflow::Output>& inputs,