Support SubSplitInfo in MultiSplitInfo.

PiperOrigin-RevId: 715420599
tensorflow · Jan 15, 2025 · bc7ec10 · bc7ec10
1 parent 1322866
commit bc7ec10
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/tensorflow_datasets/core/splits.py b/tensorflow_datasets/core/splits.py
@@ -282,9 +282,11 @@ class MultiSplitInfo(SplitInfo):
   This should only be used to read data and not when producing data.
   """
 
-  split_infos: list[SplitInfo] = dataclasses.field(default_factory=list)
+  split_infos: list[SplitInfo | SubSplitInfo] = dataclasses.field(
+      default_factory=list
+  )
 
-  def __init__(self, name: str, split_infos: list[SplitInfo]):
+  def __init__(self, name: str, split_infos: list[SplitInfo | SubSplitInfo]):
     if not split_infos:
       raise ValueError('Need to pass a non-empty list of SplitInfos')
     object.__setattr__(self, 'split_infos', split_infos)