Make the split keyword optional for load_partition

[adam-narozniak]

Issue

Currently, when you create the FederatedDataset object and want to load a partition, you need to do the following.

fds = FederatedDataset(dataset="mnist", partitioners={"train": 10})
partition = fds.load_partition(idx=0, split="train")

Yet it is clear that there is only one instance of Partitioner. Therefore, the call can be simplified. To say it differently, if when specyfing the split keyword in the load_partitition(idx=idx, split="train") train is the only possibility since you denote which partitioner will be used and only a single partitioner was used.

Proposal

Make the split keyword optional for load_partition in case of having just one partitioner:

fds = FederatedDataset(dataset="mnist", partitioners={"train": 10})
partition = fds.load_partition(idx=0) # no need to give the split here, it will be taken from "train"

Explanation

This PR does the following:

• make the split optional - code modification of FederatedDataset,
• check if it works - add tests
• (create a utility function to compare the Datasets objects).