diff --git a/grain/_src/python/dataset/transformations/packing_concat_then_split.py b/grain/_src/python/dataset/transformations/packing_concat_then_split.py index 988573072..dc48ec1e9 100644 --- a/grain/_src/python/dataset/transformations/packing_concat_then_split.py +++ b/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.""" diff --git a/grain/_src/python/dataset/transformations/packing_concat_then_split_test.py b/grain/_src/python/dataset/transformations/packing_concat_then_split_test.py index 8cfddf891..cab8dc41f 100644 --- a/grain/_src/python/dataset/transformations/packing_concat_then_split_test.py +++ b/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