Feature request: Convert standard dataset into a federated dataset

Synthetic federated datasets can constructed from standard centralized ones by artificially splitting them among clients. This is usually done using a Dirichlet distribution (e.g. [Hsu et al. 2019](https://arxiv.org/pdf/1909.06335.pdf)).
Such synthetic datasets are very useful since we can explicitly control the total number of users, as well as the heterogeneity.

It would be great to have primitives which can automatically convert standard numpy dataset into a FedJax datset.