Using comm.Gather in mpi4py with different array shapes

48 Views Asked by At

I am using mpi4py to parallise my code in python. comm.Gather gathers numpy arrays of all cores to one array automatically. I have used this and it works, but when I try to gather arrays of different shapes (which could be easily concatenated) the code fails. My code:

from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()


a = np.zeros((2 if rank==1 else 5, 3),dtype=float)+rank #size=3: shape[0] is 5,2,5 for ranks 0,1,2
#a = np.zeros((6,3),dtype=float)+rank #size=2: shape[0] is 6,6 for ranks 0,1 (this works)
print(rank,a)

b = np.zeros((12,3),dtype=float)-1
comm.Gather(a,b,root=0)
if rank==0:
    print(b)

My question now is how to write the code to concatenate the arrays of different shapes so that it works. What I don't want is to split up the arrays so that every array has the same shape and then comm.send the rest manually and concatenate manually.

0

There are 0 best solutions below