[RFC] More support for functionalities from `itertools`

🚀 The feature

Over time, we have received more and more request for additional IterDataPipe (e.g. #648, #754, plus many more). Sometimes, these functionalities are very similar to what is already implemented in itertools and more-itertools.

Keep adding more IterDataPipe one at a time seems unsustainable(?). Perhaps, we should draw a line somewhere or provide better interface for users to directly use functions from itertools. At the same time, providing APIs with names that are already familiar to Python users can improve the user experience. As @msaroufim mentioned, the Core library does aim to match operators with what is available in numpy.

We will need to decide on:

1. Coverage - which set of functionalities should we officially in torchdata?
2. Implementation - how will users be able to invoke those functions?

Coverage

0. Arbitrary based on estimated user requests/contributions
1. itertools ~20 functions (some of which already exist in torchdata)
   • This seems common enough and reasonable?
2. more-itertools ~100 functions?
   • This is probably too much.

If we provide a good wrapper, we might not need to worry about the actual coverage too much?

Implementation

0. Keep adding each function as a new IterDataPipe
   • This is what we have been doing. We can keep doing that but the cost of maintenance will increase over time.

Currently, you can use IterableWrapper, but it doesn't always work well since it accepts an iterable, and an iterable doesn't guarantee to restart if you call iter() on it again.

from torchdata.datapipes.iter import IterableWrapper
from itertools import accumulate

source_dp = IterableWrapper(range(10))
dp3 = IterableWrapper(accumulate(source_dp), deepcopy=False)
list(dp3)  # [0, 1, 3, 6, 10, 15, 21, 28, 36, 45]
list(dp3)  # []

One idea to work around that is to:

1. Provide a different wrapper that accepts a Callable that returns an Iterable, which will be iterated over

   • Users can use functool.partial to pass in arguments (including DataPipes if desired)
   • I personally think we should do this since the cost of doing so is low and unlocks other possibilities.

2. Create an Itertools DataPipe that delegates other DataPipes, it might look some like this:

class ItertoolsIterDataPipe(IterDataPipe):

    supported_operations: Dict[str, Callable] = {
        "repeat": Repeater,
        "chain": Concater,
        "filterfalse": filter_false_constructor,
        # most/all 20 `itertools` functions here?
    }

    def __new__(cls, name, *args, **kwargs):
        if name not in cls.supported_operations:
            raise RuntimeError("Operator is not supported")
        constructor = cls.supported_operations[name]
        return constructor(*args, **kwargs)

source_dp = IterableWrapper(range(10))
dp1 = source_dp.filter(lambda x: x >= 5)
dp2 = ItertoolsIterDataPipe("filterfalse", source_dp, lambda x: x >= 5)

list(dp1)  # [5, 6, 7, 8, 9]
list(dp2)  # [0, 1, 2, 3, 4]

These options are incomplete. If you have more ideas, please comment below.

Motivation, pitch

These functionalities are commonly used and can be valuable for users.

Additional context

Credit to @NicolasHug @msaroufim @pmeier and many others for past feedback and discussion related to this topic.

cc: @VitalyFedyunin @ejguan

Start by reading the existing IterDataPipe and IterableWrapper behavior, then compare the issue's itertools and more-itertools coverage options. Review how current IterDataPipes such as filter are exposed and how IterableWrapper handles repeated iteration. Done means agreeing on coverage and an invocation or wrapper design; the issue does not identify specific files or tests.