Restore ConcatThenSplitIterDataset from a JSON-serialized state

grain/_src/python/dataset/transformations/packing_concat_then_split.py

@@ -209,6 +209,13 @@ class _CtsState:
   remainder_slices: Mapping[str, tuple[int, int]]
   elements_from_buffer_after_checkpoint: int
 
+  def __post_init__(self):
+    # A state read back from JSON holds lists where the slices were tuples,
+    # and a list never equals `_EMPTY_SLICE`.
+    self.remainder_slices = {
+        key: tuple(sl) for key, sl in self.remainder_slices.items()
+    }
+
   @property
   def has_remainder(self) -> bool:
     """Whether there is a remainder element."""

grain/_src/python/dataset/transformations/packing_concat_then_split_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Unit tests for ConcatThenSplitIterDataset."""
 
+import json
 from typing import Any
 
 from absl.testing import absltest
@@ -349,9 +350,13 @@ def test_missingsequence_length(self):
   @parameterized.product(
       checkpoint_steps=[{}, {2}, {0, 3}, {0, 1, 2, 3, 4, 5}],
       split_full_length_features=[True, False],
+      through_json=[False, True],
   )
   def test_checkpointing(
-      self, checkpoint_steps: set[int], split_full_length_features: bool
+      self,
+      checkpoint_steps: set[int],
+      split_full_length_features: bool,
+      through_json: bool,
   ):
     def _create_iter(state: dict[str, Any] | None):
       ds = packing_concat_then_split.ConcatThenSplitIterDataset(
@@ -362,6 +367,8 @@ def _create_iter(state: dict[str, Any] | None):
       )
       ds_iter = ds.__iter__()
       if state is not None:
+        if through_json:
+          state = json.loads(json.dumps(state))
         ds_iter.set_state(state)
       return ds_iter