Enable more partitioners specification in more types

[adam-narozniak]

Enable specifying partitioners as type Dict[str, Union[Paritioner, int]]

How the library is currently used to create 10 partitions:

fds = FederatedDataset(dataset="mnist", partitioners={"train": 10})

Here are the new ways of specifying partitioners. All of the examples accomplish the same goal:

Dictionary-based object creation with Partitioner(s)

fds = FederatedDataset(dataset="mnist", partitioners={"train": IidPartitioner(num_partitions=10)})

Dictionary-based object creation with ints

fds = FederatedDataset(dataset="mnist", partitioners={"train":10})

Another use case is.
Dictionary-based object creation with mixed types

fds = FederatedDataset(
  dataset="mnist", 
  partitioners={
    "train":10,
    "valid": IidPartitioner(num_partitions=100)
  }
)

[danieljanes]

Lgtm!