Add ElasticIterDatasetIterator to handle scaling up and down between checkpoints.

* Allows users to keep their pipelines elastic and restore from a checkpoint with variable amount of hosts

Limitations
* Does not guarantee determinism between scaling
* The limit of parallelism is the number of shards
* Currently doesn't support multiprocessing

PiperOrigin-RevId: 864910611

Author: PAenugula

CHANGELOG.md

@@ -68,6 +68,7 @@ changes. Best viewed [here](https://google-grain.readthedocs.io/en/latest/change
   * Exposes `SharedMemoryArrayMetadata` in a public API as a metadata descriptor
     for `SharedMemoryArray`.
   * `ParquetIterDataset` can read from multiple string paths interleaving reads.
+  * Add `ElasticIterDatasetIterator` for scaling up and down the number of shards between checkpoints.
 
 * Breaking changes:
   * Custom implementations of `RandomAccessDataSource` should accept `int`

grain/_src/python/dataset/BUILD

@@ -166,6 +166,7 @@ py_library(
         ":dataset",
         "//grain/_src/core:sharding",
         "//grain/_src/python:options",
+        "//grain/_src/python/dataset/transformations:zip",
     ],
 )