Expose StreamingDataset world (or world_size and rank) as argument

[lukemelas]

Hello authors, thank you for your very nice work.

I would like to manually specify a world size and rank for a StreamingDataset (which I'm getting from a DeviceMesh). As far as I can see there is no way of doing this at the moment.

Am I missing something? If not, would you be open to a PR exposing (or expose yourself) the world in the StreamingDataset constructor (or alternatively just a world size and rank)?

All the best,
Luke