grain.experimental.assert_equal_output_after_checkpoint

grain.experimental.assert_equal_output_after_checkpoint(ds)

Tests restoring an iterator to various checkpointed states.

Parameters:

ds (Any) – The dataset to test. It is recommended to use a small dataset, potentially created using grain.python.experimental.LimitIterDataset, to restrict the number of steps being tested. The underlying dataset iterator must implement get_state and set_state for checkpointing.