Feat: WIP - Enable ND-Distributed arrays on operators

[mrava87]

Just toying around with the idea of allowing ND-Distributed arrays to be consumed by operators (like in PyLops)

A few thoughts from this initial experiment:

• self.dims/self.dimsd must be added to MPILinearOperator like we do for PyLops LinearOperator and we must ensure that they are defined for every operator (without them it is impossible to know the kind of ND input the operator would accept)
• in pylops we follow this convention:
  • matvec/rmatvec always take 1d arrays
  • dot can take ndarray (ravels them, runs matvec, reshapes them back)
  • in each operator the _matvec/_rmatvec implementations assume 1d inputs but with @reshape can reshape it for processing and then flatten it back
  • solvers can taken Nd arrays but internally the flatten, run and reshape
  • everything is a view so back and forth between 1d and nd is rather cheap
• in pylops-mpi, any ravel/reshape does inherently create a new Distributed array and does copy local arrays so it is kind of expensive. We may need to consider to change a bit the places where ravel/reshape is done, perhaps:
  • in matvec/rmatvec of MPILinearOperator instead of in dot, so that also in solvers we can just call those methods on NDarrays
  • avoid ravel before calling _matvec/_rmatvec and instead change @reshape to just pass blank if the input is already ND (maybe with a check that its global_shape is aligned with dims/dimsd... and similar avoid flattening it back

@rohanbabbar04 FYI, this is as far as I got, some tentative code changes, an example, and some thoughts that came to my mind whilst toying around 😄

[rohanbabbar04]

@mrava87, thanks this makes sense, looks like a good start.

I will add some sub-task pointers in a couple of days here itself, which we would need to cover 🙂 .