Added default no-op function for flatmap

test/test_iterdatapipe.py

@@ -697,6 +697,7 @@ def fn(e):
 
         flatmapped_dp = source_dp.flatmap(fn)
         expected_list = list(itertools.chain(*[(e, e * 10) for e in source_dp]))
+
         self.assertEqual(expected_list, list(flatmapped_dp))
 
         # Funtional Test: Specify input_col
@@ -713,6 +714,19 @@ def mul_fn(a, b):
         input_col_2_dp = tuple_source_dp.flatmap(mul_fn, input_col=(0, 2))
         self.assertEqual(list(itertools.chain(*[(-2, 2) for _ in range(20)])), list(input_col_2_dp))
 
+        # flatmap with no fn specified
+        default_dp = tuple_source_dp.flatmap()
+        self.assertEqual(list(itertools.chain(*[(n - 1, n, n + 1) for n in range(20)])), list(default_dp))
+
+        # flatmap with no fn specified, multiple input_col
+        default_dp = tuple_source_dp.flatmap(input_col=(0, 2))
+        self.assertEqual(list(itertools.chain(*[(n - 1, n + 1) for n in range(20)])), list(default_dp))
+
+        # flatmap with no fn specified, some special input
+        tuple_source_dp = IterableWrapper([[1, 2, [3, 4]], [5, 6, [7, 8]]])
+        default_dp = tuple_source_dp.flatmap(input_col=(0, 2))
+        self.assertEqual([1, [3, 4], 5, [7, 8]], list(default_dp))
+
         # Reset Test: reset the DataPipe after reading part of it
         n_elements_before_reset = 5
         res_before_reset, res_after_reset = reset_after_n_next_calls(flatmapped_dp, n_elements_before_reset)

torchdata/datapipes/iter/transform/callable.py

@@ -13,6 +13,15 @@
 T_co = TypeVar("T_co", covariant=True)
 
 
+def _no_op_fn(*args):
+    """
+    No-operation function, returns passed arguments.
+    """
+    if len(args) == 1:
+        return args[0]
+    return args
+
+
 @functional_datapipe("map_batches")
 class BatchMapperIterDataPipe(IterDataPipe[T_co]):
     r"""
@@ -99,6 +108,7 @@ class FlatMapperIterDataPipe(IterDataPipe[T_co]):
 
     Note:
         The output from ``fn`` must be a Sequence. Otherwise, an error will be raised.
+        If ``fn`` is ``None``, source DataPipe will be just flattened vertically, provided that items can be unpacked.
 
     Args:
         datapipe: Source IterDataPipe
@@ -112,13 +122,20 @@ class FlatMapperIterDataPipe(IterDataPipe[T_co]):
         >>> flatmapped_dp = source_dp.flatmap(fn)
         >>> list(flatmapped_dp)
         [0, 0, 1, 10, 2, 20, 3, 30, 4, 40]
+        >>>
+        >>> source_dp = IterableWrapper([[1, 2, 3], [4, 5, 6]])
+        >>> flatmapped_dp = source_dp.flatmap()
+        >>> list(flatmapped_dp)
+        [1, 2, 3, 4, 5, 6]
     """
     datapipe: IterDataPipe
     fn: Callable
 
-    def __init__(self, datapipe: IterDataPipe, fn: Callable, input_col=None) -> None:
+    def __init__(self, datapipe: IterDataPipe, fn: Callable = None, input_col=None) -> None:
         self.datapipe = datapipe
 
+        if fn is None:
+            fn = _no_op_fn
         _check_unpickable_fn(fn)
         self.fn = fn  # type: ignore[assignment]
         self.input_col = input_col