TF_RETURN_IF_ERROR(GraphToFunctionDef(*g, new_name, &replace_fdef));
    TF_RETURN_IF_ERROR(fld->AddFunctionDef(replace_fdef));
    loop_cond_func->set_name(new_name);
    while_node->ClearAttr("cond");
    while_node->AddAttr("cond", *loop_cond_func);
  }

  return absl::OkStatus();
}

// Rewrites while loop cond function for host.
Status RewriteHostWhileLoopCond(

    // expected-remark@above {{Predecessors: {4}}}
  }
  func.return
  // expected-remark@above {{ID: 7}}
  // expected-remark@above {{Sinks: {6}}}
}

// CHECK-LABEL: func @while_body
func.func @while_body(%arg0: tensor<i1>) -> tensor<i1> {
  // expected-remark@above {{ID: 5}}
  %graph = tf_executor.graph {
  // expected-remark@above {{ID: 3}}
    %island:2 = tf_executor.island {

    variable ops, thus allowing them to run in parallel across iterations.

    Before:

    ```mlir
    !tf_res = type tensor<!tf_type.resource<tensor<f32>>>
    func @while_body(%arg0: !tf_res, %arg1: !tf_res, %arg2: tensor<f32>, %arg3: tensor<f32>) -> (!tf_res, !tf_res, tensor<f32>, tensor<f32>) {
      %graph:4 = tf_executor.graph {