Repository: tensorflow
Issue: tf.train.Checkpoint.restore does not restore the tf.data.Iterator state.
Labels: stat:awaiting tensorflower,type:bug,comp:data,TF 2.13

### Issue type

Bug

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

No

### Source

binary

### TensorFlow version

tf 2.13.0

### Custom code

Yes

### OS platform and distribution

Ubuntu 22.04

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

11.8

### GPU model and memory

_No response_

### Current behavior?

`tf.data.Dataset` iterators are supported in `tf.train.Checkpoint`.  When I try to restore a checkpoint, I always get a `StopIteration` error.  This should not be the case because : 

1. I save the checkpoint only at the end of an epoch.
2. Before saving the checkpoint, I always call `db_iter=iter(dataset)` ,so it should give me a new iterator.

In prior versions of TF there were functions like `make_initializable_iterator()` which are now deprecated. Since, there is no native mechanism to reset a python iterator to beginning, I wonder why `tf.train.Checkpoint.restore()` does not save the iterator state properly.

I also understand well that one can just loop over the dataset like:

```
import tensorflow as tf
ds = tf.data.Dataset.range(50)
for index, data in ds:
    # Do something 
```

However, saving the dataset in the checkpoint has two problems:

a) **I have not explicitly tested this**  For a very large dataset like MSCOCO or OpenImages or even Imagenet, what does saving the dataset mean , if I already have the dataset as TFRecords ? Will it end up saving the whole dataset again inside the  checkpoint ? At least, `tf.data.Dataset.save()` points towards this. In case this is it, I would never like to save a large dataset directly into the checkpoint.

b) **This is fully tested** If I experiment with a small simple dataset ( e.g:- `tf.data.Dataset.range(10)` ) and save it in the checkpoint, then it does not restore the state of the iterator. This observation is the same as #48178 which is still unresolved. 

### Standalone code to reproduce the issue

```shell
import os
import numpy as np
import tensorflow as tf
from absl import app


class Net(tf.keras.Model):
    def __init__(self):
        super(Net, self).__init__()
        self.x = tf.Variable(tf.keras.initializers.GlorotUniform()(shape=[1]))

    def call(self, inputs):
        return inputs


def main(argv):
    del argv
    tf.random.set_seed(123)
    net = Net()
    optim = tf.optimizers.Adam(learning_rate=0.001)

    dataset = tf.data.Dataset.from_tensor_slices((np.arange(10)))
    dataset = dataset.shuffle(10, reshuffle_each_iteration=True)
    dataset = dataset.batch(2)
    db_iter = iter(dataset)
    step = tf.Variable(0)
    checkpoint = tf.train.Checkpoint(
        step=step, optimizer=optim, net=net, db_iter=db_iter
    )
    manager = tf.train.CheckpointManager(checkpoint, "./ckpts", max_to_keep=50)
    manager.restore_or_initialize()
    print(step)
    for epoch in range(10):
        manager.restore_or_initialize()
        for _ in range(dataset.cardinality()):
            batch = next(db_iter)
            step.assign_add(1)
            print(batch)
        db_iter = iter(dataset)
        manager.save()
        print(f"Epoch {epoch} finished.")


if __name__ == "__main__":
    app.run(main)
```
```


### Relevant log output

```shell
The first time, it will print something like : 

<tf.Variable 'Variable:0' shape=() dtype=int32, numpy=0>
tf.Tensor([3 6], shape=(2,), dtype=int64)
tf.Tensor([4 0], shape=(2,), dtype=int64)
tf.Tensor([1 5], shape=(2,), dtype=int64)
tf.Tensor([7 8], shape=(2,), dtype=int64)
tf.Tensor([2 9], shape=(2,), dtype=int64)
Epoch 0 finished.
tf.Tensor([6 3], shape=(2,), dtype=int64)
tf.Tensor([0 5], shape=(2,), dtype=int64)
tf.Tensor([2 8], shape=(2,), dtype=int64)
tf.Tensor([1 7], shape=(2,), dtype=int64)
tf.Tensor([4 9], shape=(2,), dtype=int64)
Epoch 1 finished.
tf.Tensor([0 3], shape=(2,), dtype=int64)
tf.Tensor([1 4], shape=(2,), dtype=int64)
tf.Tensor([6 7], shape=(2,), dtype=int64)
tf.Tensor([9 8], shape=(2,), dtype=int64)
tf.Tensor([5 2], shape=(2,), dtype=int64)
Epoch 2 finished.
tf.Tensor([3 7], shape=(2,), dtype=int64)
tf.Tensor([0 9], shape=(2,), dtype=int64)
tf.Tensor([4 2], shape=(2,), dtype=int64)
tf.Tensor([1 8], shape=(2,), dtype=int64)
tf.Tensor([5 6], shape=(2,), dtype=int64)

and so on....


From next time onwards:

    raise StopIteration
StopIteration
```