feat: added masked option to asarray

[mrava87]

This PR adds a new optional input to the DistributedArray.asarray method called masked, which can be used with distributed arrays that have a mask to create a global view of the distributed array that contains only elements from one color of the mask (so that repeated portions of the array are not returned).

[mrava87]

@astroC86, this is relevant to your PR as we can now compare the full output of your MPIMatrixMult with that of the serial MatrixMult.

@tharittk I have made some changes to the MPI tests, it would be good to do the same for the NCCL ones and check that the code I added works also with that backend (it should but I have not tested it yet)

[rohanbabbar04]

Looks good to me!!

[tharittk]

I have tested it by copy-paste your update. This uncovers the bug in nccl_asarray() function when working with subcomm. I am not sure exactly where yet but it looks like current asarray() will call nccl_asarray() with self.local_shapes (which call all GPUs in the world communicator) and may mess with the subcomm.

The MPI does not have this problem because it sends the object and no buffer / array unrolling needed.

[mrava87]

I have tested it by copy-paste your update. This uncovers the bug in nccl_asarray() function when working with subcomm. I am not sure exactly where yet but it looks like current asarray() will call nccl_asarray() with self.local_shapes (which call all GPUs in the world communicator) and may mess with the subcomm.

The MPI does not have this problem because it sends the object and no buffer / array unrolling needed.

@tharittk thanks for checking! I am not sure I understand what the problem is with NCCL but I agree it may be due to self.local_shapes that is not in the subcomm but in the total comm... I suggest that I merge this as is and make a Github Issue pointing to the problem with NCCL so you can handle it in a separate PR adding also tests and @astroC86 can use these additions in his PR (which targets only MPI so far)