Register Functional API to different DataPipe class

🚀 The feature

@functional_datapipe("func")
class IterToMapMapDataPipe(MapDataPipe):
    def __init__(self, datapipe: IterDataPipe):
        self.datapipe = datapipe

Using the functional_datapipe above would register this functional API to MapDataPipe. As we want to register this functional_datapiep to the IterDataPipe, I have to do the following trick.

class IterToMapMapDataPipe(MapDataPipe):
    def __init__(self, datapipe: IterDataPipe):
        self.datapipe = datapipe

IterDataPipe.register_datapipe_as_function("func", IterToMapMapDataPipe)

Motivation, pitch

Option 1:
Use the type annotation of the first input datapipe from __init__ or __new__ function to determine which DataPipe class we should register.

class IterToMapMapDataPipe(MapDataPipe):
    def __init__(self, datapipe: IterDataPipe):  # <---------- IterDataPipe here is used to register
        self.datapipe = datapipe

Option 2:
Add an optional argument for functional_datapipe API to explicitly set the right DataPipe class.

@functional_datapipe("func", register_cls=IterDataPipe)
class IterToMapMapDataPipe(MapDataPipe):
    def __init__(self, datapipe: IterDataPipe):
        self.datapipe = datapipe

Alternatives

No response

Additional context

No response

Start by reading the functional_datapipe API and the register_datapipe_as_function entry point mentioned in the issue. Compare the proposed annotation-based and explicit register_cls approaches, then determine which registration behavior the project should support. Done means a functional API can be registered against a different DataPipe class without the manual workaround.