IRMapping values_in_scope;
  // Function arguments can appear as operands, so they clone should
  // be aware of them.
  assert(subgraph.FuncArguments().size() == new_func.getNumArguments());
  for (int i = 0; i < subgraph.FuncArguments().size(); ++i) {
    Value original_value = subgraph.FuncArguments()[i];
    Value new_func_arg = new_func.getArgument(i);
    values_in_scope.map(original_value, new_func_arg);
  }

            func::FuncOp func = while_op.body_function();
            auto func_argument = func.getArgument(next_op_operand_num);
            // Check if the op is returned without mutation. Returning values
            // from a while op follow return or identity -> return pattern.
            if (checkIfAnyUserIsConnectedToTermiantor(func_argument))
              next_values_to_visit.push_back(