Repository: pytorch
Issue: Add support for `__collate__` attrib on dataset elements in `default_collate`
Labels: module: dataloader,triaged

### 🚀 The feature, motivation and pitch

There seemed to have been some discussions on alternative ways to supply collate function (like in [#33181](https://github.com/pytorch/pytorch/issues/33181)). I also found is slightly inconvenient to add custom way to collate metadata on my dataset. My pitch to solve this would be to add check if the dataset element has `__collate__` attrib and use it if present.

This would help to easily extend the behavior of `default_collate`. Simple example:
```python
class MyMetadata(dict):
    @staticmethod
    def __collate__(batch, *, collate_fn_map: Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]] = None):
        # when collating metadata (basically dictionaries), we want to return batch as a simple list
        return batch

class MyDataset(torch.utils.data.Dataset):
    def __len__(self):
        return 16
    
    def __getitem__(self, index):
        x = torch.tensor(index)
        y = torch.tensor(index)
        metadata = {"index": index, "list_of_variable_length": list(range(index))}
        return {"x": x, "y": y, "metadata": MyMetadata(metadata)}


ds = MyDataset()
dl = torch.utils.data.DataLoader(ds, batch_size=4)

next(iter(dl))
```

Expected output would be:
```
{'x': tensor([0, 1, 2, 3]),
 'y': tensor([0, 1, 2, 3]),
 'metadata': [{'index': 0, 'list_of_variable_length': []},
  {'index': 1, 'list_of_variable_length': [0]},
  {'index': 2, 'list_of_variable_length': [0, 1]},
  {'index': 3, 'list_of_variable_length': [0, 1, 2]}]}
```

This seems to be easy to implement (although it it my first time contributing here, so I don't know for sure). Something like:

```python
def pitched_collate(batch, *, collate_fn_map: Optional[Dict[Union[Type, Tuple[Type, ...]], Callable]] = None):
    elem = batch[0]

    #HERE ARE THE ONLY CHANGES COMPARED TO collate
    if hasattr(elem, "__collate__"):
        return getattr(elem, "__collate__")(batch, collate_fn_map=collate_fn_map)
     #END OF CHANGES

    elem_type = type(elem)

    if collate_fn_map is not None:
        if elem_type in collate_fn_map:
            return collate_fn_map[elem_type](batch, collate_fn_map=collate_fn_map)

        for collate_type in collate_fn_map:
            if isinstance(elem, collate_type):
                return collate_fn_map[collate_type](batch, collate_fn_map=collate_fn_map)

    if isinstance(elem, collections.abc.Mapping):
        try:
            return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
        except TypeError:
            # The mapping type may not support `__init__(iterable)`.
            return {key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem}
    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
        return elem_type(*(collate(samples, collate_fn_map=collate_fn_map) for samples in zip(*batch)))
    elif isinstance(elem, collections.abc.Sequence):
        # check to make sure that the elements in batch have consistent size
        it = iter(batch)
        elem_size = len(next(it))
        if not all(len(elem) == elem_size for elem in it):
            raise RuntimeError('each element in list of batch should be of equal size')
        transposed = list(zip(*batch))  # It may be accessed twice, so we use a list.

        if isinstance(elem, tuple):
            return [collate(samples, collate_fn_map=collate_fn_map) for samples in transposed]  # Backwards compatibility.
        else:
            try:
                return elem_type([collate(samples, collate_fn_map=collate_fn_map) for samples in transposed])
            except TypeError:
                # The sequence type may not support `__init__(iterable)` (e.g., `range`).
                return [collate(samples, collate_fn_map=collate_fn_map) for samples in transposed]

    raise TypeError(default_collate_err_msg_format.format(elem_type))

```

### Alternatives
Perhaps an option to register new type in `default_collate_fn_map` would solve this, but from what it seems there is currently no way to even access `default_collate_fn_map` without importing it from `torch.utils.data._utils.collate`

### Additional context

If this seems like a good addition, I'd be willing to try and add this in a PR.

cc @SsnL @VitalyFedyunin @ejguan @NivekT @dzhulgakov