SmallPtrSet<Operation*, 4> ops_connected_to_fetch, int num_old_outputs) {
  assert(num_old_outputs + resource_equivalence_classes.getNumClasses() ==
         func.getNumArguments());
  auto graph_op = cast<GraphOp>(func.front().front());

  auto fetch = graph_op.GetFetch();
  OpBuilder builder_chain_src(fetch);
  builder_chain_src.setInsertionPointToStart(fetch->getBlock());

        *module_op->getBodyRegion().getOps().begin());
    if (!func_op) {
      return absl::InternalError("Input MLIR must have only 1 func");
    }
    if (arguments.size() != func_op.getNumArguments()) {
      return absl::InternalError("Input argument has wrong size");
    }

    // Convert input xla::Literal arguments to constants, this allows using
    // constant folding to evaluate function return value.

        decomposed_partitioned_call_callees) {
  OpBuilder builder(while_op);
  auto modify_region_arguments = [&](Region& region) {
    int64_t original_arg_count = region.getNumArguments();
    for (int64_t i = 0; i < original_arg_count; ++i) {
      auto operand = while_op.getOperand(i);
      auto it = buffer_to_size->find(operand);
      if (it == buffer_to_size->end()) continue;