Is there a way to map a TensorDict?

[NightMachinery]

I am currently using HuggingFace Datasets to load, process, and save some data. However, HF Datasets saves the data in the Arrow format and wastes a lot of time converting between Arrow and PyTorch tensors.

I am wondering if I can use memory-mapped TensorDicts for this purpose?

How can I do a map on batches of TensorDict?

Looking through the tutorials, the nearest example I found was using a Dataloader with collate_fn as the map function:

DataLoader(
    tensor_dict, batch_size=batch_size, collate_fn=map_fn,
)

But this does not allow me to form a pipeline of map functions. I also don't know how to save and load the resulting dataset.

[vmoens]

This is on our radar!
I like the idea of having a (possibly multiprocessed) map to execute some transform over all the elements of a tensordict.
Stay tuned, I'll ping you once we have PR with this