}

      argument.replaceAllUsesWith(it->second.getOutput());
      arg_indexes_to_remove.set(argument.getArgNumber());
    }

    // Erase function arguments with bounded input.
    func.eraseArguments(arg_indexes_to_remove);
  }
};

}  // namespace

absl::Status AddSessionInitializerAndInlineCheckpoint(