Repository: tensorflow
Issue: Checkpointing support for GPU datasets
Labels: stat:awaiting tensorflower,type:feature,comp:apis,comp:gpu

### Issue type

Feature Request

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

binary

### TensorFlow version

tf 2.15

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 22.04

### Mobile device

_No response_

### Python version

3.11

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

12.3

### GPU model and memory

_No response_

### Current behavior?

Hello! Saving a checkpoint for a dataset placed on the GPU fails with:
```
Cannot assign a device for operation SerializeIterator (...) because no supported kernel for GPU devices is available.
```
(full error message below)

I expected the checkpoint for dataset to be saved as it would be for CPU-placed dataset. 

Is there a way to overcome this limitation? Is it possible for you to add checkpointing support for GPU datasets?

Thank you for your help!

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import tensorflow_datasets as tfds

ds = tfds.load('mnist', split='train', shuffle_files=True)
ds = ds.apply(tf.data.experimental.prefetch_to_device("/gpu:0"))

it = iter(ds)
mgr = tf.train.Checkpoint(it)
cpt = mgr.save("/tmp/")
```


### Relevant log output

```shell
Traceback (most recent call last):
  File "/home/skarpinski/tfgpu.py", line 9, in <module>
    cpt = mgr.save("/tmp/")
          ^^^^^^^^^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/checkpoint/checkpoint.py", line 2533, in save
    return self._write(
           ^^^^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/checkpoint/checkpoint.py", line 2369, in _write
    output = self._saver.save(file_prefix=file_prefix, options=options)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/checkpoint/checkpoint.py", line 1356, in save
    save_path, new_feed_additions = self._save_cached_when_graph_building(
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/checkpoint/checkpoint.py", line 1284, in _save_cached_when_graph_building
    self._gather_serialized_tensors(object_graph_tensor))
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/checkpoint/checkpoint.py", line 1245, in _gather_serialized_tensors
    save_util.serialize_graph_view(self._graph_view,
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/checkpoint/save_util.py", line 309, in serialize_graph_view
    serialized_tensors = _get_and_write_tensors_to_serialize(
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/checkpoint/save_util.py", line 157, in _get_and_write_tensors_to_serialize
    trackable, tensor_dict = _get_tensors_from_legacy_saveable(
                             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/checkpoint/save_util.py", line 190, in _get_tensors_from_legacy_saveable
    save_util_v1.generate_saveable_objects(
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/checkpoint/save_util_v1.py", line 180, in generate_saveable_objects
    maybe_saveable = saveable_object_util.create_saveable_object(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/training/saving/saveable_object_util.py", line 473, in create_saveable_object
    return factory(name=key)
           ^^^^^^^^^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/training/saving/saveable_object_util.py", line 536, in create_saveable
    tensor_dict = save_fn()
                  ^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 900, in _serialize_to_tensors
    serialized_iterator = gen_dataset_ops.serialize_iterator(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/ops/gen_dataset_ops.py", line 6523, in serialize_iterator
    _ops.raise_from_not_ok_status(e, name)
  File "/home/skarpinski/.local/lib/python3.11/site-packages/tensorflow/python/framework/ops.py", line 5983, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot assign a device for operation SerializeIterator: Could not satisfy explicit device specification '/job:localhost/replica:0/task:0/device:GPU:0' because no supported kernel for GPU devices is available.
Colocation Debug Info:
Colocation group had the following types and supported devices: 
Root Member(assigned_device_name_index_=1 requested_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' assigned_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' resource_device_name_='/job:localhost/replica:0/task:0/device:GPU:0' supported_device_types_=[CPU] possible_devices_=[]
SerializeIterator: CPU 
_Arg: GPU CPU 

Colocation members, user-requested devices, and framework assigned devices, if any:
  resource__handle (_Arg)  framework assigned device=/job:localhost/replica:0/task:0/device:GPU:0
  SerializeIterator (SerializeIterator) /job:localhost/replica:0/task:0/device:GPU:0

Op: SerializeIterator
Node attrs: external_state_policy=2
Registered kernels:
  device='CPU'

         [[{{node SerializeIterator}}]] [Op:SerializeIterator] name:
```