Repository: keras-team/keras

Issue:
`split_dataset` cannot split a dataset containing nested tuples
I'm using tensorflow and loading data into a `tf.data.Dataset` object structured like so, where each dataset has the same length (20160 entries in this case):
```python
data1 = tf.data.Dataset(...) # tensorspec of [1, 16384, 2]
data2 = tf.data.Dataset(...) # tensorspec of [512, 32, 2]
labels = tf.data.Dataset(...) # tensorspec of scalar int32

nested = (data1, data2)
dataset = tf.data.Dataset.zip(nested, labels)
```
This is to comply with the input to [`model.fit`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit):
> A [tf.data.Dataset](https://www.tensorflow.org/api_docs/python/tf/data/Dataset). Should return a tuple of either (inputs, targets) or (inputs, targets, sample_weights)

However, when trying to split into training and validation sets, `_restore_dataset_from_list` throws an error. Specifically, this line:
https://github.com/keras-team/keras/blob/26f438edd5516bf02504a5a0e13a20ceec69cb29/keras/src/utils/dataset_utils.py#L455

It's hard for me to see what's happening exactly because `np.array` is what's throwing the error:
```
setting an array element with a sequence. The requested array has an inhomogeneous shape after 2 dimensions.
The detected shape was (20160, 2) + inhomogeneous part
```

But this seems to indicate the inhomogenous part is the label, and if it's trying to coerce the structure I think it is (below) then it would  indeed be a violation of the allowed structure for a numpy array
```python
[
    [[d1, d2], label]
    ...
]
```
Given that a `tf.data.Dataset` explicitly supports nesting and `model.fit` expects a tuple of inputs for a multi-input model, I think `split_dataset` should be able to handle nesting as well. But if handling nesting requires a slower method, then the current method should be kept and nesting handled as a fallback in case `np.array` fails.

My workaround for the time being is to zip a flat dataset and then pack the inputs *after* the call to `split_dataset` using `map`:
```python
dataset = tf.data.Dataset.zip(data1, data2, labels)
trainer, valer = tf.keras.utils.split_dataset(dataset, 0.7)
trainer = trainer.map(lambda a,b,l: ((a,b), l))
valer   =   valer.map(lambda a,b,l: ((a,b), l))
```
Which isn't awful, but it took me a very long time to figure out the issue. I'd just like to save others the headache :)

Edit: or does the label dataset need to be a list of 1D tensors with single elements instead of scalars?


Discussion:
Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSdHag0RVFS7UXzZkKcsFCKOcX8raCupKK9RHSlYxp5U8lSJbQ/viewform?entry.492125872=Yes&entry.243948740=https://github.com/keras-team/keras/issues/19797">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSdHag0RVFS7UXzZkKcsFCKOcX8raCupKK9RHSlYxp5U8lSJbQ/viewform?entry.492125872=No&entry.243948740=https://github.com/keras-team/keras/issues/19797">No</a>

Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSdHag0RVFS7UXzZkKcsFCKOcX8raCupKK9RHSlYxp5U8lSJbQ/viewform?entry.492125872=Yes&entry.243948740=https://github.com/keras-team/keras/issues/19797">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSdHag0RVFS7UXzZkKcsFCKOcX8raCupKK9RHSlYxp5U8lSJbQ/viewform?entry.492125872=No&entry.243948740=https://github.com/keras-team/keras/issues/19797">No</a>




Patch:
diff --git a/keras/src/utils/dataset_utils.py b/keras/src/utils/dataset_utils.py
index ff27552278f4..ba36818c48e0 100644
--- a/keras/src/utils/dataset_utils.py
+++ b/keras/src/utils/dataset_utils.py
@@ -6,6 +6,7 @@
 
 import numpy as np
 
+from keras.src import tree
 from keras.src.api_export import keras_export
 from keras.src.utils import io_utils
 from keras.src.utils.module_utils import tensorflow as tf
@@ -137,16 +138,7 @@ def _convert_dataset_to_list(
         data_size_warning_flag,
         start_time,
     ):
-        if dataset_type_spec in [tuple, list]:
-            # The try-except here is for NumPy 1.24 compatibility, see:
-            # https://numpy.org/neps/nep-0034-infer-dtype-is-object.html
-            try:
-                arr = np.array(sample)
-            except ValueError:
-                arr = np.array(sample, dtype=object)
-            dataset_as_list.append(arr)
-        else:
-            dataset_as_list.append(sample)
+        dataset_as_list.append(sample)
 
     return dataset_as_list
 
@@ -169,23 +161,23 @@ def _get_data_iterator_from_dataset(dataset, dataset_type_spec):
                 "Please provide a non-empty list of arrays."
             )
 
-        if _get_type_spec(dataset[0]) is np.ndarray:
-            expected_shape = dataset[0].shape
-            for i, element in enumerate(dataset):
-                if np.array(element).shape[0] != expected_shape[0]:
-                    raise ValueError(
-                        "Received a list of NumPy arrays with different "
-                        f"lengths. Mismatch found at index {i}, "
-                        f"Expected shape={expected_shape} "
-                        f"Received shape={np.array(element).shape}."
-                        "Please provide a list of NumPy arrays with "
-                        "the same length."
-                    )
-        else:
-            raise ValueError(
-                "Expected a list of `numpy.ndarray` objects,"
-                f"Received: {type(dataset[0])}"
-            )
+        expected_shape = None
+        for i, element in enumerate(dataset):
+            if not isinstance(element, np.ndarray):
+                raise ValueError(
+                    "Expected a list of `numpy.ndarray` objects,"
+                    f"Received: {type(element)} at index {i}."
+                )
+            if expected_shape is None:
+                expected_shape = element.shape
+            elif element.shape[0] != expected_shape[0]:
+                raise ValueError(
+                    "Received a list of NumPy arrays with different lengths."
+                    f"Mismatch found at index {i}, "
+                    f"Expected shape={expected_shape} "
+                    f"Received shape={np.array(element).shape}."
+                    "Please provide a list of NumPy arrays of the same length."
+                )
 
         return iter(zip(*dataset))
     elif dataset_type_spec == tuple:
@@ -195,23 +187,23 @@ def _get_data_iterator_from_dataset(dataset, dataset_type_spec):
                 "Please provide a non-empty tuple of arrays."
             )
 
-        if _get_type_spec(dataset[0]) is np.ndarray:
-            expected_shape = dataset[0].shape
-            for i, element in enumerate(dataset):
-                if np.array(element).shape[0] != expected_shape[0]:
-                    raise ValueError(
-                        "Received a tuple of NumPy arrays with different "
-                        f"lengths. Mismatch found at index {i}, "
-                        f"Expected shape={expected_shape} "
-                        f"Received shape={np.array(element).shape}."
-                        "Please provide a tuple of NumPy arrays with "
-                        "the same length."
-                    )
-        else:
-            raise ValueError(
-                "Expected a tuple of `numpy.ndarray` objects, "
-                f"Received: {type(dataset[0])}"
-            )
+        expected_shape = None
+        for i, element in enumerate(dataset):
+            if not isinstance(element, np.ndarray):
+                raise ValueError(
+                    "Expected a tuple of `numpy.ndarray` objects,"
+                    f"Received: {type(element)} at index {i}."
+                )
+            if expected_shape is None:
+                expected_shape = element.shape
+            elif element.shape[0] != expected_shape[0]:
+                raise ValueError(
+                    "Received a tuple of NumPy arrays with different lengths."
+                    f"Mismatch found at index {i}, "
+                    f"Expected shape={expected_shape} "
+                    f"Received shape={np.array(element).shape}."
+                    "Please provide a tuple of NumPy arrays of the same length."
+                )
 
         return iter(zip(*dataset))
     elif dataset_type_spec == tf.data.Dataset:
@@ -436,23 +428,24 @@ def _restore_dataset_from_list(
     dataset_as_list, dataset_type_spec, original_dataset
 ):
     """Restore the dataset from the list of arrays."""
-    if dataset_type_spec in [tuple, list]:
-        return tuple(np.array(sample) for sample in zip(*dataset_as_list))
-    elif dataset_type_spec == tf.data.Dataset:
-        if isinstance(original_dataset.element_spec, dict):
-            restored_dataset = {}
-            for d in dataset_as_list:
-                for k, v in d.items():
-                    if k not in restored_dataset:
-                        restored_dataset[k] = [v]
-                    else:
-                        restored_dataset[k].append(v)
-            return restored_dataset
-        else:
-            return tuple(np.array(sample) for sample in zip(*dataset_as_list))
+    if dataset_type_spec in [tuple, list, tf.data.Dataset] or is_torch_dataset(
+        original_dataset
+    ):
+        # Save structure by taking the first element.
+        element_spec = dataset_as_list[0]
+        # Flatten each element.
+        dataset_as_list = [tree.flatten(sample) for sample in dataset_as_list]
+        # Combine respective elements at all indices.
+        dataset_as_list = [np.array(sample) for sample in zip(*dataset_as_list)]
+        # Recreate the original structure of elements.
+        dataset_as_list = tree.pack_sequence_as(element_spec, dataset_as_list)
+        # Turn lists to tuples as tf.data will fail on lists.
+        return tree.traverse(
+            lambda x: tuple(x) if isinstance(x, list) else x,
+            dataset_as_list,
+            top_down=False,
+        )
 
-    elif is_torch_dataset(original_dataset):
-        return tuple(np.array(sample) for sample in zip(*dataset_as_list))
     return dataset_as_list
 
 
@@ -477,14 +470,12 @@ def _get_type_spec(dataset):
         return list
     elif isinstance(dataset, np.ndarray):
         return np.ndarray
-    elif isinstance(dataset, dict):
-        return dict
     elif isinstance(dataset, tf.data.Dataset):
         return tf.data.Dataset
     elif is_torch_dataset(dataset):
-        from torch.utils.data import Dataset as torchDataset
+        from torch.utils.data import Dataset as TorchDataset
 
-        return torchDataset
+        return TorchDataset
     else:
         return None
 


Test patch:
diff --git a/keras/src/utils/dataset_utils_test.py b/keras/src/utils/dataset_utils_test.py
index 11cb275ff815..b87e571d94b6 100644
--- a/keras/src/utils/dataset_utils_test.py
+++ b/keras/src/utils/dataset_utils_test.py
@@ -1,145 +1,49 @@
+import itertools
+
 import numpy as np
+from absl.testing import parameterized
+from torch.utils.data import Dataset as TorchDataset
 
 from keras.src.testing import test_case
+from keras.src.testing.test_utils import named_product
 from keras.src.utils.dataset_utils import split_dataset
 from keras.src.utils.module_utils import tensorflow as tf
 
 
-class DatasetUtilsTest(test_case.TestCase):
-    def test_split_dataset_list(self):
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        dataset = [
-            np.random.sample((n_sample, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        ]
-        dataset_left, dataset_right = split_dataset(
-            dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            int(dataset_left.cardinality()), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            int(dataset_right.cardinality()), int(n_sample * right_size)
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (n_cols)
-        )
+class MyTorchDataset(TorchDataset):
 
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        dataset = [
-            np.random.sample((n_sample, 100, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        ]
-        dataset_left, dataset_right = split_dataset(
-            dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            int(dataset_left.cardinality()), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            int(dataset_right.cardinality()), int(n_sample * right_size)
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (100, n_cols)
-        )
+    def __init__(self, x, y):
+        self.x = x
+        self.y = y
 
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        dataset = [
-            np.random.sample((n_sample, 10, 10, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        ]
-        dataset_left, dataset_right = split_dataset(
-            dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            int(dataset_left.cardinality()), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            int(dataset_right.cardinality()), int(n_sample * right_size)
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (10, 10, n_cols)
-        )
+    def __len__(self):
+        return len(self.x)
 
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        dataset = [
-            np.random.sample((n_sample, 100, 10, 30, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        ]
-        dataset_left, dataset_right = split_dataset(
-            dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            int(dataset_left.cardinality()), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            int(dataset_right.cardinality()), int(n_sample * right_size)
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape,
-            (100, 10, 30, n_cols),
-        )
+    def __getitem__(self, index):
+        return self.x[index], self.y[index]
 
-    def test_split_dataset_tuple(self):
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        dataset = (
-            np.random.sample((n_sample, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        dataset_left, dataset_right = split_dataset(
-            dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            int(dataset_left.cardinality()), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            int(dataset_right.cardinality()), int(n_sample * right_size)
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (n_cols)
-        )
 
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        dataset = (
-            np.random.sample((n_sample, 100, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        dataset_left, dataset_right = split_dataset(
-            dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            int(dataset_left.cardinality()), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            int(dataset_right.cardinality()), int(n_sample * right_size)
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (100, n_cols)
+class DatasetUtilsTest(test_case.TestCase, parameterized.TestCase):
+    @parameterized.named_parameters(
+        named_product(
+            dataset_type=["list", "tuple", "tensorflow", "torch"],
+            features_shape=[(2,), (100, 2), (10, 10, 2)],
         )
+    )
+    def test_split_dataset(self, dataset_type, features_shape):
+        n_sample, left_size, right_size = 100, 0.2, 0.8
+        features = np.random.sample((n_sample,) + features_shape)
+        labels = np.random.sample((n_sample, 1))
 
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        dataset = (
-            np.random.sample((n_sample, 10, 10, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        dataset_left, dataset_right = split_dataset(
-            dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            int(dataset_left.cardinality()), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            int(dataset_right.cardinality()), int(n_sample * right_size)
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (10, 10, n_cols)
-        )
+        if dataset_type == "list":
+            dataset = [features, labels]
+        elif dataset_type == "tuple":
+            dataset = (features, labels)
+        elif dataset_type == "tensorflow":
+            dataset = tf.data.Dataset.from_tensor_slices((features, labels))
+        elif dataset_type == "torch":
+            dataset = MyTorchDataset(features, labels)
 
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        dataset = (
-            np.random.sample((n_sample, 100, 10, 30, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
         dataset_left, dataset_right = split_dataset(
             dataset, left_size=left_size, right_size=right_size
         )
@@ -149,77 +53,30 @@ def test_split_dataset_tuple(self):
         self.assertEqual(
             int(dataset_right.cardinality()), int(n_sample * right_size)
         )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape,
-            (100, 10, 30, n_cols),
-        )
-
-    def test_split_dataset_tensorflow(self):
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        features, labels = (
-            np.random.sample((n_sample, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        tf_dataset = tf.data.Dataset.from_tensor_slices((features, labels))
-        dataset_left, dataset_right = split_dataset(
-            tf_dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            int(dataset_left.cardinality()), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            int(dataset_right.cardinality()), int(n_sample * right_size)
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (n_cols)
-        )
-
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        features, labels = (
-            np.random.sample((n_sample, 100, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        tf_dataset = tf.data.Dataset.from_tensor_slices((features, labels))
-        dataset_left, dataset_right = split_dataset(
-            tf_dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            int(dataset_left.cardinality()), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            int(dataset_right.cardinality()), int(n_sample * right_size)
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (100, n_cols)
-        )
+        for sample in itertools.chain(dataset_left, dataset_right):
+            self.assertEqual(sample[0].shape, features_shape)
+            self.assertEqual(sample[1].shape, (1,))
+
+    @parameterized.named_parameters(
+        named_product(structure_type=["dict", "tuple"])
+    )
+    def test_split_dataset_nested_structures(self, structure_type):
+        n_sample, left_size, right_size = 100, 0.2, 0.8
+        features1 = np.random.sample((n_sample, 2))
+        features2 = np.random.sample((n_sample, 10, 2))
+        labels = np.random.sample((n_sample, 1))
+
+        if structure_type == "dict":
+            dataset = tf.data.Dataset.from_tensor_slices(
+                {"x1": features1, "x2": features2, "labels": labels}
+            )
+        elif structure_type == "tuple":
+            dataset = tf.data.Dataset.from_tensor_slices(
+                ((features1, features2), labels)
+            )
 
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        features, labels = (
-            np.random.sample((n_sample, 10, 10, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        tf_dataset = tf.data.Dataset.from_tensor_slices((features, labels))
         dataset_left, dataset_right = split_dataset(
-            tf_dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            int(dataset_left.cardinality()), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            int(dataset_right.cardinality()), int(n_sample * right_size)
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (10, 10, n_cols)
-        )
-
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        features, labels = (
-            np.random.sample((n_sample, 100, 10, 30, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        tf_dataset = tf.data.Dataset.from_tensor_slices((features, labels))
-        dataset_left, dataset_right = split_dataset(
-            tf_dataset, left_size=left_size, right_size=right_size
+            dataset, left_size=left_size, right_size=right_size
         )
         self.assertEqual(
             int(dataset_left.cardinality()), int(n_sample * left_size)
@@ -227,108 +84,11 @@ def test_split_dataset_tensorflow(self):
         self.assertEqual(
             int(dataset_right.cardinality()), int(n_sample * right_size)
         )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape,
-            (100, 10, 30, n_cols),
-        )
-
-    def test_split_dataset_torch(self):
-        # sample torch dataset class
-        from torch.utils.data import Dataset as torchDataset
-
-        class Dataset(torchDataset):
-            "Characterizes a dataset for PyTorch"
-
-            def __init__(self, x, y):
-                "Initialization"
-                self.x = x
-                self.y = y
-
-            def __len__(self):
-                "Denotes the total number of samples"
-                return len(self.x)
-
-            def __getitem__(self, index):
-                "Generates one sample of data"
-                return self.x[index], self.y[index]
-
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        features, labels = (
-            np.random.sample((n_sample, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        torch_dataset = Dataset(features, labels)
-        dataset_left, dataset_right = split_dataset(
-            torch_dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            len([sample for sample in dataset_left]), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            len([sample for sample in dataset_right]),
-            int(n_sample * right_size),
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (n_cols,)
-        )
-
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        features, labels = (
-            np.random.sample((n_sample, 100, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        torch_dataset = Dataset(features, labels)
-        dataset_left, dataset_right = split_dataset(
-            torch_dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            len([sample for sample in dataset_left]), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            len([sample for sample in dataset_right]),
-            int(n_sample * right_size),
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (100, n_cols)
-        )
-
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        features, labels = (
-            np.random.sample((n_sample, 10, 10, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        torch_dataset = Dataset(features, labels)
-        dataset_left, dataset_right = split_dataset(
-            torch_dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            len([sample for sample in dataset_left]), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            len([sample for sample in dataset_right]),
-            int(n_sample * right_size),
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape, (10, 10, n_cols)
-        )
-
-        n_sample, n_cols, n_pred, left_size, right_size = 100, 2, 1, 0.2, 0.8
-        features, labels = (
-            np.random.sample((n_sample, 100, 10, 30, n_cols)),
-            np.random.sample((n_sample, n_pred)),
-        )
-        torch_dataset = Dataset(features, labels)
-        dataset_left, dataset_right = split_dataset(
-            torch_dataset, left_size=left_size, right_size=right_size
-        )
-        self.assertEqual(
-            len([sample for sample in dataset_left]), int(n_sample * left_size)
-        )
-        self.assertEqual(
-            len([sample for sample in dataset_right]),
-            int(n_sample * right_size),
-        )
-        self.assertEqual(
-            [sample for sample in dataset_right][0][0].shape,
-            (100, 10, 30, n_cols),
-        )
+        for sample in itertools.chain(dataset_left, dataset_right):
+            if structure_type == "dict":
+                x1, x2, labels = sample["x1"], sample["x2"], sample["labels"]
+            elif structure_type == "tuple":
+                (x1, x2), labels = sample
+            self.assertEqual(x1.shape, (2,))
+            self.assertEqual(x2.shape, (10, 2))
+            self.assertEqual(labels.shape, (1,))