Access to key in DatasetDict map
Indeed that would be cool

Also FYI right now the easiest way to do this is
```python
dataset_dict["train"] = dataset_dict["train"].map(my_transform_for_the_train_set)
dataset_dict["test"] = dataset_dict["test"].map(my_transform_for_the_test_set)
```
It is possible that we want to do different things in the `map` function (and possibly other functions too) of a `DatasetDict`, depending on the key. I understand that `DatasetDict.map` is a really thin wrapper of `Dataset.map`, so it is easy to directly implement this functionality in the client code. Still, it'd be nice if there can be a flag, similar to `with_indices`, that allows the callable to know the key inside `DatasetDict`.