Allow pickling support for sub classes of DatasetBuilder.

PiperOrigin-RevId: 681701404

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/dataset_builder.py

@@ -354,10 +354,13 @@ def code_path(cls) -> Optional[epath.Path]:
       return epath.Path(filepath)
 
   def __getstate__(self):
-    return self._original_state
+    # This needs to be of the same format as __dict__,
+    # to match with the pickling implementation of the derived classes.
+    return dict(_original_state=self._original_state)
 
   def __setstate__(self, state):
-    self.__init__(**state)
+    original_state = state["_original_state"]
+    self.__init__(**original_state)
 
   @functools.cached_property
   def canonical_version(self) -> utils.Version: