Finish getting TfdsDataSource mock to iterate correctly.

The numpy_function has to return 1 output, so stack serialized outputs. This allows the lookup inside TfDataLoader to batch normally.

PiperOrigin-RevId: 628080888

Author: The TensorFlow Datasets Authors

tensorflow_datasets/testing/mocking.py

@@ -116,10 +116,13 @@ def _getitems(
     serialized: bool = False,
 ) -> Sequence[Any]:
   """Function to overwrite __getitems__ in data sources."""
-  return [
+  items = [
       _getitem(self, record_key, generator, serialized=serialized)
       for record_key in record_keys
   ]
+  if serialized:
+    return np.array(items)
+  return items
 
 
 def _deserialize_example_np(serialized_example, *, decoders=None):