File size: 32,334 Bytes
d1ceb73
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
# This base template ("datapipe.pyi.in") is generated from mypy stubgen with minimal editing for code injection
# The output file will be "datapipe.pyi". This is executed as part of torch/CMakeLists.txt
# Note that, for mypy, .pyi file takes precedent over .py file, such that we must define the interface for other
# classes/objects here, even though we are not injecting extra code into them at the moment.

from typing import Any, Callable, Dict, Generic, Iterator, List, Literal, Optional, TypeVar, Union

from torch.utils.data import Dataset, default_collate, IterableDataset
from torch.utils.data.datapipes._hook_iterator import _SnapshotState
from torch.utils.data.datapipes._typing import _DataPipeMeta, _IterDataPipeMeta

T_co = TypeVar("T_co", covariant=True)
T = TypeVar("T")
UNTRACABLE_DATAFRAME_PIPES: Any

class MapDataPipe(Dataset[T_co], metaclass=_DataPipeMeta):
    functions: Dict[str, Callable] = ...
    reduce_ex_hook: Optional[Callable] = ...
    getstate_hook: Optional[Callable] = ...
    str_hook: Optional[Callable] = ...
    repr_hook: Optional[Callable] = ...
    def __getattr__(self, attribute_name: Any): ...
    @classmethod
    def register_function(cls, function_name: Any, function: Any) -> None: ...
    @classmethod
    def register_datapipe_as_function(
        cls,
        function_name: Any,
        cls_to_register: Any,
    ): ...
    def __getstate__(self): ...
    def __reduce_ex__(self, *args: Any, **kwargs: Any): ...
    @classmethod
    def set_getstate_hook(cls, hook_fn: Any) -> None: ...
    @classmethod
    def set_reduce_ex_hook(cls, hook_fn: Any) -> None: ...
    # Functional form of 'BatcherMapDataPipe'
    def batch(self, batch_size: int, drop_last: bool = False, wrapper_class=DataChunk) -> MapDataPipe:
        r"""
        Create mini-batches of data (functional name: ``batch``).
    
        An outer dimension will be added as ``batch_size`` if ``drop_last`` is set to ``True``,
        or ``length % batch_size`` for the last batch if ``drop_last`` is set to ``False``.
    
        Args:
            datapipe: Iterable DataPipe being batched
            batch_size: The size of each batch
            drop_last: Option to drop the last batch if it's not full
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> dp = SequenceWrapper(range(10))
            >>> batch_dp = dp.batch(batch_size=2)
            >>> list(batch_dp)
            [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]
        """
    
    # Functional form of 'ConcaterMapDataPipe'
    def concat(self, *datapipes: MapDataPipe) -> MapDataPipe:
        r"""
        Concatenate multiple Map DataPipes (functional name: ``concat``).
    
        The new index of is the cumulative sum of source DataPipes.
        For example, if there are 2 source DataPipes both with length 5,
        index 0 to 4 of the resulting `ConcatMapDataPipe` would refer to
        elements of the first DataPipe, and 5 to 9 would refer to elements
        of the second DataPipe.
    
        Args:
            datapipes: Map DataPipes being concatenated
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> dp1 = SequenceWrapper(range(3))
            >>> dp2 = SequenceWrapper(range(3))
            >>> concat_dp = dp1.concat(dp2)
            >>> list(concat_dp)
            [0, 1, 2, 0, 1, 2]
        """
    
    # Functional form of 'MapperMapDataPipe'
    def map(self, fn: Callable= ...) -> MapDataPipe:
        r"""
        Apply the input function over each item from the source DataPipe (functional name: ``map``).
    
        The function can be any regular Python function or partial object. Lambda
        function is not recommended as it is not supported by pickle.
    
        Args:
            datapipe: Source MapDataPipe
            fn: Function being applied to each item
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.map import SequenceWrapper, Mapper
            >>> def add_one(x):
            ...     return x + 1
            >>> dp = SequenceWrapper(range(10))
            >>> map_dp_1 = dp.map(add_one)
            >>> list(map_dp_1)
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
            >>> map_dp_2 = Mapper(dp, lambda x: x + 1)
            >>> list(map_dp_2)
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        """
    
    # Functional form of 'ShufflerIterDataPipe'
    def shuffle(self, *, indices: Optional[List] = None) -> IterDataPipe:
        r"""
        Shuffle the input MapDataPipe via its indices (functional name: ``shuffle``).
    
        When it is used with :class:`~torch.utils.data.DataLoader`, the methods to
        set up random seed are different based on :attr:`num_workers`.
    
        For single-process mode (:attr:`num_workers == 0`), the random seed is set before
        the :class:`~torch.utils.data.DataLoader` in the main process. For multi-process
        mode (:attr:`num_worker > 0`), ``worker_init_fn`` is used to set up a random seed
        for each worker process.
    
        Args:
            datapipe: MapDataPipe being shuffled
            indices: a list of indices of the MapDataPipe. If not provided, we assume it uses 0-based indexing
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> dp = SequenceWrapper(range(10))
            >>> shuffle_dp = dp.shuffle().set_seed(0)
            >>> list(shuffle_dp)
            [7, 8, 1, 5, 3, 4, 2, 0, 9, 6]
            >>> list(shuffle_dp)
            [6, 1, 9, 5, 2, 4, 7, 3, 8, 0]
            >>> # Reset seed for Shuffler
            >>> shuffle_dp = shuffle_dp.set_seed(0)
            >>> list(shuffle_dp)
            [7, 8, 1, 5, 3, 4, 2, 0, 9, 6]
    
        Note:
            Even thought this ``shuffle`` operation takes a ``MapDataPipe`` as the input, it would return an
            ``IterDataPipe`` rather than a ``MapDataPipe``, because ``MapDataPipe`` should be non-sensitive to
            the order of data order for the sake of random reads, but ``IterDataPipe`` depends on the order
            of data during data-processing.
        """
    
    # Functional form of 'ZipperMapDataPipe'
    def zip(self, *datapipes: MapDataPipe[T_co]) -> MapDataPipe:
        r"""
        Aggregates elements into a tuple from each of the input DataPipes (functional name: ``zip``).
    
        This MataPipe is out of bound as soon as the shortest input DataPipe is exhausted.
    
        Args:
            *datapipes: Map DataPipes being aggregated
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.map import SequenceWrapper
            >>> dp1 = SequenceWrapper(range(3))
            >>> dp2 = SequenceWrapper(range(10, 13))
            >>> zip_dp = dp1.zip(dp2)
            >>> list(zip_dp)
            [(0, 10), (1, 11), (2, 12)]
        """
    

class IterDataPipe(IterableDataset[T_co], metaclass=_IterDataPipeMeta):
    functions: Dict[str, Callable] = ...
    reduce_ex_hook: Optional[Callable] = ...
    getstate_hook: Optional[Callable] = ...
    str_hook: Optional[Callable] = ...
    repr_hook: Optional[Callable] = ...
    _number_of_samples_yielded: int = ...
    _snapshot_state: _SnapshotState = _SnapshotState.Iterating
    _fast_forward_iterator: Optional[Iterator] = ...
    def __getattr__(self, attribute_name: Any): ...
    @classmethod
    def register_function(cls, function_name: Any, function: Any) -> None: ...
    @classmethod
    def register_datapipe_as_function(
        cls,
        function_name: Any,
        cls_to_register: Any,
        enable_df_api_tracing: bool = ...,
    ): ...
    def __getstate__(self): ...
    def __reduce_ex__(self, *args: Any, **kwargs: Any): ...
    @classmethod
    def set_getstate_hook(cls, hook_fn: Any) -> None: ...
    @classmethod
    def set_reduce_ex_hook(cls, hook_fn: Any) -> None: ...
    # Functional form of 'BatcherIterDataPipe'
    def batch(self, batch_size: int, drop_last: bool = False, wrapper_class=DataChunk) -> IterDataPipe:
        r"""
        Creates mini-batches of data (functional name: ``batch``).
    
        An outer dimension will be added as ``batch_size`` if ``drop_last`` is set to ``True``, or ``length % batch_size`` for the
        last batch if ``drop_last`` is set to ``False``.
    
        Args:
            datapipe: Iterable DataPipe being batched
            batch_size: The size of each batch
            drop_last: Option to drop the last batch if it's not full
            wrapper_class: wrapper to apply onto each batch (type ``List``) before yielding,
                defaults to ``DataChunk``
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.iter import IterableWrapper
            >>> dp = IterableWrapper(range(10))
            >>> dp = dp.batch(batch_size=3, drop_last=True)
            >>> list(dp)
            [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
        """
    
    # Functional form of 'CollatorIterDataPipe'
    def collate(self, conversion: Optional[Union[Callable[..., Any],Dict[Union[str, Any], Union[Callable, Any]],]] = default_collate, collate_fn: Optional[Callable] = None) -> IterDataPipe:
        r"""
        Collates samples from DataPipe to Tensor(s) by a custom collate function (functional name: ``collate``).
    
        By default, it uses :func:`torch.utils.data.default_collate`.
    
        .. note::
            While writing a custom collate function, you can import :func:`torch.utils.data.default_collate` for the
            default behavior and `functools.partial` to specify any additional arguments.
    
        Args:
            datapipe: Iterable DataPipe being collated
            collate_fn: Customized collate function to collect and combine data or a batch of data.
                Default function collates to Tensor(s) based on data type.
    
        Example:
            >>> # xdoctest: +SKIP
            >>> # Convert integer data to float Tensor
            >>> class MyIterDataPipe(torch.utils.data.IterDataPipe):
            ...     def __init__(self, start, end):
            ...         super(MyIterDataPipe).__init__()
            ...         assert end > start, "this example code only works with end >= start"
            ...         self.start = start
            ...         self.end = end
            ...
            ...     def __iter__(self):
            ...         return iter(range(self.start, self.end))
            ...
            ...     def __len__(self):
            ...         return self.end - self.start
            ...
            >>> ds = MyIterDataPipe(start=3, end=7)
            >>> print(list(ds))
            [3, 4, 5, 6]
            >>> def collate_fn(batch):
            ...     return torch.tensor(batch, dtype=torch.float)
            ...
            >>> collated_ds = CollateIterDataPipe(ds, collate_fn=collate_fn)
            >>> print(list(collated_ds))
            [tensor(3.), tensor(4.), tensor(5.), tensor(6.)]
        """
    
    # Functional form of 'ConcaterIterDataPipe'
    def concat(self, *datapipes: IterDataPipe) -> IterDataPipe:
        r"""
        Concatenates multiple Iterable DataPipes (functional name: ``concat``).
    
        The resulting DataPipe will yield all the elements from the first input DataPipe, before yielding from the subsequent ones.
    
        Args:
            datapipes: Iterable DataPipes being concatenated
    
        Example:
            >>> # xdoctest: +REQUIRES(module:torchdata)
            >>> import random
            >>> from torchdata.datapipes.iter import IterableWrapper
            >>> dp1 = IterableWrapper(range(3))
            >>> dp2 = IterableWrapper(range(5))
            >>> list(dp1.concat(dp2))
            [0, 1, 2, 0, 1, 2, 3, 4]
        """
    
    # Functional form of 'DemultiplexerIterDataPipe'
    def demux(self, num_instances: int, classifier_fn: Callable[[T_co], Optional[int]], drop_none: bool = False, buffer_size: int = 1000) -> List[IterDataPipe]:
        r"""
        Splits the input DataPipe into multiple child DataPipes, using the given classification function (functional name: ``demux``).
    
        A list of the child DataPipes is returned from this operation.
    
        Args:
            datapipe: Iterable DataPipe being filtered
            num_instances: number of instances of the DataPipe to create
            classifier_fn: a function that maps values to an integer within the range ``[0, num_instances - 1]`` or ``None``
            drop_none: defaults to ``False``, if ``True``, the function will skip over elements classified as ``None``
            buffer_size: this defines the maximum number of inputs that the buffer can hold across all child
                DataPipes while waiting for their values to be yielded.
                Defaults to ``1000``. Use ``-1`` for the unlimited buffer.
    
        Examples:
            >>> # xdoctest: +REQUIRES(module:torchdata)
            >>> from torchdata.datapipes.iter import IterableWrapper
            >>> def odd_or_even(n):
            ...     return n % 2
            >>> source_dp = IterableWrapper(range(5))
            >>> dp1, dp2 = source_dp.demux(num_instances=2, classifier_fn=odd_or_even)
            >>> list(dp1)
            [0, 2, 4]
            >>> list(dp2)
            [1, 3]
            >>> # It can also filter out any element that gets `None` from the `classifier_fn`
            >>> def odd_or_even_no_zero(n):
            ...     return n % 2 if n != 0 else None
            >>> dp1, dp2 = source_dp.demux(num_instances=2, classifier_fn=odd_or_even_no_zero, drop_none=True)
            >>> list(dp1)
            [2, 4]
            >>> list(dp2)
            [1, 3]
        """
    
    # Functional form of 'FilterIterDataPipe'
    def filter(self, filter_fn: Callable, input_col=None) -> IterDataPipe:
        r"""
        Filters out elements from the source datapipe according to input ``filter_fn`` (functional name: ``filter``).
    
        Args:
            datapipe: Iterable DataPipe being filtered
            filter_fn: Customized function mapping an element to a boolean.
            input_col: Index or indices of data which ``filter_fn`` is applied, such as:
    
                - ``None`` as default to apply ``filter_fn`` to the data directly.
                - Integer(s) is used for list/tuple.
                - Key(s) is used for dict.
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.iter import IterableWrapper
            >>> def is_even(n):
            ...     return n % 2 == 0
            >>> dp = IterableWrapper(range(5))
            >>> filter_dp = dp.filter(filter_fn=is_even)
            >>> list(filter_dp)
            [0, 2, 4]
        """
    
    # Functional form of 'ForkerIterDataPipe'
    def fork(self, num_instances: int, buffer_size: int = 1000, copy: Optional[Literal["shallow", "deep"]] = None) -> List[IterDataPipe]:
        r"""
        Creates multiple instances of the same Iterable DataPipe (functional name: ``fork``).
    
        Args:
            datapipe: Iterable DataPipe being copied
            num_instances: number of instances of the datapipe to create
            buffer_size: this restricts how far ahead the leading child DataPipe
               can read relative to the slowest child DataPipe.
               Defaults to ``1000``. Use ``-1`` for the unlimited buffer.
            copy: copy strategy to use for items yielded by each branch. Supported
                options are ``None`` for no copying, ``"shallow"`` for shallow object
                copies, and ``"deep"`` for deep object copies. Defaults to ``None``.
    
        Note:
            All branches of the forked pipeline return the identical object unless
            the copy parameter is supplied. If the object is mutable or contains
            mutable objects, changing them in one branch will affect all others.
    
        Example:
            >>> # xdoctest: +REQUIRES(module:torchdata)
            >>> from torchdata.datapipes.iter import IterableWrapper
            >>> source_dp = IterableWrapper(range(5))
            >>> dp1, dp2 = source_dp.fork(num_instances=2)
            >>> list(dp1)
            [0, 1, 2, 3, 4]
            >>> list(dp2)
            [0, 1, 2, 3, 4]
        """
    
    # Functional form of 'GrouperIterDataPipe'
    def groupby(self, group_key_fn: Callable[[T_co], Any], *, keep_key: bool = False, buffer_size: int = 10000, group_size: Optional[int] = None, guaranteed_group_size: Optional[int] = None, drop_remaining: bool = False) -> IterDataPipe:
        r"""
        Groups data from IterDataPipe by keys from ``group_key_fn``, yielding a ``DataChunk`` with batch size up to ``group_size``.
    
        (functional name: ``groupby``).
    
        The samples are read sequentially from the source ``datapipe``, and a batch of samples belonging to the same group
        will be yielded as soon as the size of the batch reaches ``group_size``. When the buffer is full,
        the DataPipe will yield the largest batch with the same key, provided that its size is larger
        than ``guaranteed_group_size``. If its size is smaller, it will be dropped if ``drop_remaining=True``.
    
        After iterating through the entirety of source ``datapipe``, everything not dropped due to the buffer capacity
        will be yielded from the buffer, even if the group sizes are smaller than ``guaranteed_group_size``.
    
        Args:
            datapipe: Iterable datapipe to be grouped
            group_key_fn: Function used to generate group key from the data of the source datapipe
            keep_key: Option to yield the matching key along with the items in a tuple,
                resulting in `(key, [items])` otherwise returning [items]
            buffer_size: The size of buffer for ungrouped data
            group_size: The max size of each group, a batch is yielded as soon as it reaches this size
            guaranteed_group_size: The guaranteed minimum group size to be yielded in case the buffer is full
            drop_remaining: Specifies if the group smaller than ``guaranteed_group_size`` will be dropped from buffer
                when the buffer is full
    
        Example:
            >>> import os
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.iter import IterableWrapper
            >>> def group_fn(file):
            ...     return os.path.basename(file).split(".")[0]
            >>> source_dp = IterableWrapper(["a.png", "b.png", "a.json", "b.json", "a.jpg", "c.json"])
            >>> dp0 = source_dp.groupby(group_key_fn=group_fn)
            >>> list(dp0)
            [['a.png', 'a.json', 'a.jpg'], ['b.png', 'b.json'], ['c.json']]
            >>> # A group is yielded as soon as its size equals to `group_size`
            >>> dp1 = source_dp.groupby(group_key_fn=group_fn, group_size=2)
            >>> list(dp1)
            [['a.png', 'a.json'], ['b.png', 'b.json'], ['a.jpg'], ['c.json']]
            >>> # Scenario where `buffer` is full, and group 'a' needs to be yielded since its size > `guaranteed_group_size`
            >>> dp2 = source_dp.groupby(group_key_fn=group_fn, buffer_size=3, group_size=3, guaranteed_group_size=2)
            >>> list(dp2)
            [['a.png', 'a.json'], ['b.png', 'b.json'], ['a.jpg'], ['c.json']]
        """
    
    # Functional form of 'FileListerIterDataPipe'
    def list_files(self, masks: Union[str, List[str]] = '', *, recursive: bool = False, abspath: bool = False, non_deterministic: bool = False, length: int = -1) -> IterDataPipe:
        r"""
        Given path(s) to the root directory, yields file pathname(s) (path + filename) of files within the root directory.
    
        Multiple root directories can be provided (functional name: ``list_files``).
    
        Args:
            root: Root directory or a sequence of root directories
            masks: Unix style filter string or string list for filtering file name(s)
            recursive: Whether to return pathname from nested directories or not
            abspath: Whether to return relative pathname or absolute pathname
            non_deterministic: Whether to return pathname in sorted order or not.
                If ``False``, the results yielded from each root directory will be sorted
            length: Nominal length of the datapipe
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.iter import FileLister
            >>> dp = FileLister(root=".", recursive=True)
            >>> list(dp)
            ['example.py', './data/data.tar']
        """
    
    # Functional form of 'MapperIterDataPipe'
    def map(self, fn: Callable, input_col=None, output_col=None) -> IterDataPipe:
        r"""
        Applies a function over each item from the source DataPipe (functional name: ``map``).
    
        The function can be any regular Python function or partial object. Lambda
        function is not recommended as it is not supported by pickle.
    
        Args:
            datapipe: Source Iterable DataPipe
            fn: Function being applied over each item
            input_col: Index or indices of data which ``fn`` is applied, such as:
    
                - ``None`` as default to apply ``fn`` to the data directly.
                - Integer(s) is used for list/tuple.
                - Key(s) is used for dict.
    
            output_col: Index of data where result of ``fn`` is placed. ``output_col`` can be specified
                only when ``input_col`` is not ``None``
    
                - ``None`` as default to replace the index that ``input_col`` specified; For ``input_col`` with
                  multiple indices, the left-most one is used, and other indices will be removed.
                - Integer is used for list/tuple. ``-1`` represents to append result at the end.
                - Key is used for dict. New key is acceptable.
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.iter import IterableWrapper, Mapper
            >>> def add_one(x):
            ...     return x + 1
            >>> dp = IterableWrapper(range(10))
            >>> map_dp_1 = dp.map(add_one)  # Invocation via functional form is preferred
            >>> list(map_dp_1)
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
            >>> # We discourage the usage of `lambda` functions as they are not serializable with `pickle`
            >>> # Use `functools.partial` or explicitly define the function instead
            >>> map_dp_2 = Mapper(dp, lambda x: x + 1)
            >>> list(map_dp_2)
            [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
        """
    
    # Functional form of 'MultiplexerIterDataPipe'
    def mux(self, *datapipes) -> IterDataPipe:
        r"""
        Yields one element at a time from each of the input Iterable DataPipes (functional name: ``mux``).
    
        As in, one element from the 1st input DataPipe, then one element from the 2nd DataPipe in the next iteration,
        and so on. It ends when the shortest input DataPipe is exhausted.
    
        Args:
            datapipes: Iterable DataPipes that will take turn to yield their elements, until the shortest DataPipe is exhausted
    
        Example:
            >>> # xdoctest: +REQUIRES(module:torchdata)
            >>> from torchdata.datapipes.iter import IterableWrapper
            >>> dp1, dp2, dp3 = IterableWrapper(range(3)), IterableWrapper(range(10, 15)), IterableWrapper(range(20, 25))
            >>> list(dp1.mux(dp2, dp3))
            [0, 10, 20, 1, 11, 21, 2, 12, 22]
        """
    
    # Functional form of 'FileOpenerIterDataPipe'
    def open_files(self, mode: str = 'r', encoding: Optional[str] = None, length: int = -1) -> IterDataPipe:
        r"""
        Given pathnames, opens files and yield pathname and file stream in a tuple (functional name: ``open_files``).
    
        Args:
            datapipe: Iterable datapipe that provides pathnames
            mode: An optional string that specifies the mode in which
                the file is opened by ``open()``. It defaults to ``r``, other options are
                ``b`` for reading in binary mode and ``t`` for text mode.
            encoding: An optional string that specifies the encoding of the
                underlying file. It defaults to ``None`` to match the default encoding of ``open``.
            length: Nominal length of the datapipe
    
        Note:
            The opened file handles will be closed by Python's GC periodically. Users can choose
            to close them explicitly.
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.iter import FileLister, FileOpener, StreamReader
            >>> dp = FileLister(root=".").filter(lambda fname: fname.endswith('.txt'))
            >>> dp = FileOpener(dp)
            >>> dp = StreamReader(dp)
            >>> list(dp)
            [('./abc.txt', 'abc')]
        """
    
    # Functional form of 'StreamReaderIterDataPipe'
    def read_from_stream(self, chunk=None) -> IterDataPipe:
        r"""
        Given IO streams and their label names, yield bytes with label name as tuple.
    
        (functional name: ``read_from_stream``).
    
        Args:
            datapipe: Iterable DataPipe provides label/URL and byte stream
            chunk: Number of bytes to be read from stream per iteration.
                If ``None``, all bytes will be read until the EOF.
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.iter import IterableWrapper, StreamReader
            >>> from io import StringIO
            >>> dp = IterableWrapper([("alphabet", StringIO("abcde"))])
            >>> list(StreamReader(dp, chunk=1))
            [('alphabet', 'a'), ('alphabet', 'b'), ('alphabet', 'c'), ('alphabet', 'd'), ('alphabet', 'e')]
        """
    
    # Functional form of 'RoutedDecoderIterDataPipe'
    def routed_decode(self, *handlers: Callable, key_fn: Callable= ...) -> IterDataPipe:
        r"""
        Decodes binary streams from input DataPipe, yields pathname and decoded data in a tuple.
    
        (functional name: ``routed_decode``)
    
        Args:
            datapipe: Iterable datapipe that provides pathname and binary stream in tuples
            handlers: Optional user defined decoder handlers. If ``None``, basic and image decoder
                handlers will be set as default. If multiple handles are provided, the priority
                order follows the order of handlers (the first handler has the top priority)
            key_fn: Function for decoder to extract key from pathname to dispatch handlers.
                Default is set to extract file extension from pathname
    
        Note:
            When ``key_fn`` is specified returning anything other than extension, the default
            handler will not work and users need to specify custom handler. Custom handler
            could use regex to determine the eligibility to handle data.
        """
    
    # Functional form of 'ShardingFilterIterDataPipe'
    def sharding_filter(self, sharding_group_filter=None) -> IterDataPipe:
        r"""
        Wrapper that allows DataPipe to be sharded (functional name: ``sharding_filter``).
    
        After ``apply_sharding`` is called, each instance of the DataPipe (on different workers) will have every `n`-th element of the
        original DataPipe, where `n` equals to the number of instances.
    
        Args:
            source_datapipe: Iterable DataPipe that will be sharded
        """
    
    # Functional form of 'ShufflerIterDataPipe'
    def shuffle(self, *, buffer_size: int = 10000, unbatch_level: int = 0) -> IterDataPipe:
        r"""
        Shuffle the input DataPipe with a buffer (functional name: ``shuffle``).
    
        The buffer with ``buffer_size`` is filled with elements from the datapipe first. Then,
        each item will be yielded from the buffer by reservoir sampling via iterator.
    
        ``buffer_size`` is required to be larger than ``0``. For ``buffer_size == 1``, the
        datapipe is not shuffled. In order to fully shuffle all elements from datapipe,
        ``buffer_size`` is required to be greater than or equal to the size of datapipe.
    
        When it is used with :class:`torch.utils.data.DataLoader`, the methods to
        set up random seed are different based on :attr:`num_workers`.
    
        For single-process mode (:attr:`num_workers == 0`), the random seed is set before
        the :class:`~torch.utils.data.DataLoader` in the main process. For multi-process
        mode (:attr:`num_worker > 0`), `worker_init_fn` is used to set up a random seed
        for each worker process.
    
        Args:
            datapipe: The IterDataPipe being shuffled
            buffer_size: The buffer size for shuffling (default to ``10000``)
            unbatch_level: Specifies if it is necessary to unbatch source data before
                applying the shuffle
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.iter import IterableWrapper
            >>> dp = IterableWrapper(range(10))
            >>> shuffle_dp = dp.shuffle()
            >>> list(shuffle_dp)
            [0, 4, 1, 6, 3, 2, 9, 5, 7, 8]
        """
    
    # Functional form of 'UnBatcherIterDataPipe'
    def unbatch(self, unbatch_level: int = 1) -> IterDataPipe:
        r"""
        Undos batching of data (functional name: ``unbatch``).
    
        In other words, it flattens the data up to the specified level within a batched DataPipe.
    
        Args:
            datapipe: Iterable DataPipe being un-batched
            unbatch_level: Defaults to ``1`` (only flattening the top level). If set to ``2``,
                it will flatten the top two levels, and ``-1`` will flatten the entire DataPipe.
    
        Example:
            >>> # xdoctest: +SKIP
            >>> from torchdata.datapipes.iter import IterableWrapper
            >>> source_dp = IterableWrapper([[[0, 1], [2]], [[3, 4], [5]], [[6]]])
            >>> dp1 = source_dp.unbatch()
            >>> list(dp1)
            [[0, 1], [2], [3, 4], [5], [6]]
            >>> dp2 = source_dp.unbatch(unbatch_level=2)
            >>> list(dp2)
            [0, 1, 2, 3, 4, 5, 6]
        """
    
    # Functional form of 'ZipperIterDataPipe'
    def zip(self, *datapipes: IterDataPipe) -> IterDataPipe:
        r"""
        Aggregates elements into a tuple from each of the input DataPipes (functional name: ``zip``).
    
        The output is stopped as soon as the shortest input DataPipe is exhausted.
    
        Args:
            *datapipes: Iterable DataPipes being aggregated
    
        Example:
            >>> # xdoctest: +REQUIRES(module:torchdata)
            >>> from torchdata.datapipes.iter import IterableWrapper
            >>> dp1, dp2, dp3 = IterableWrapper(range(5)), IterableWrapper(range(10, 15)), IterableWrapper(range(20, 25))
            >>> list(dp1.zip(dp2, dp3))
            [(0, 10, 20), (1, 11, 21), (2, 12, 22), (3, 13, 23), (4, 14, 24)]
        """
    

class DFIterDataPipe(IterDataPipe):
    def _is_dfpipe(self): ...
    def __iter__(self): ...

class _DataPipeSerializationWrapper:
    def __init__(self, datapipe): ...
    def __getstate__(self): ...
    def __setstate__(self, state): ...
    def __len__(self): ...

class _IterDataPipeSerializationWrapper(_DataPipeSerializationWrapper, IterDataPipe):
    def __iter__(self): ...

class _MapDataPipeSerializationWrapper(_DataPipeSerializationWrapper, MapDataPipe):
    def __getitem__(self, idx): ...

class DataChunk(list, Generic[T]):
    def __init__(self, items):
        super().__init__(items)
        self.items = items
    def as_str(self, indent: str = "") -> str:
        res = indent + "[" + ", ".join(str(i) for i in iter(self)) + "]"
        return res
    def __iter__(self) -> Iterator[T]:
        yield from super().__iter__()
    def raw_iterator(self) -> T:  # type: ignore[misc]
        yield from self.items