Add local_shape to DistributedArray

[mrava87]

Motivation

The class DistributedArray currently requires users to provide the global shape of the array (via global_shape), however it does not allow users to choose how such an array is split between ranks: local_shape is found by the local_split method to ensure best load balancing between ranks. There are however scenarios where users may want to choose local_shape, for example when they know they array will be re-interpreted as nd-array by some operator and the split would need to take into account the fact dimension zero should be split whilst the other dimensions should be on the same rank.

As such, I suggest to add local_shape as new input parameter for DistributedArray.

Tasks

• Add local_shape=None as new input parameter for DistributedArray. If provided, force such splitting if not provided default to what we currently have.
• Add test and example in plot_distributed_array.py