MPI - отправка и получение столбцов матрицы

Я пытаюсь отправить столбцы матрицы другим процессам, используя Scatter. Приведенный ниже код отлично работает для строк, поэтому для отправки столбцов с минимальными изменениями я использую функцию транспонирования Numpy. Однако это, похоже, не имеет никакого эффекта, если я не сделаю полную новую копию матрицы (что, как вы можете себе представить, противоречит цели).

3 минимальных примера для иллюстрации проблемы (должны работать с 3 процессами!) ниже.

1. Разброс строк (работает как положено):

   comm = MPI.COMM_WORLD
   rank = comm.Get_rank()
   size = comm.Get_size()
   
   A = np.zeros((3,3))
   if rank==0:
       A = np.matrix([[1.,2.,3.],[4.,5.,6.],[7.,8.,9.]])
   
   local_a = np.zeros(3)
   
   comm.Scatter(A, local_a, root=0)
   print "process", rank, "has", local_a
   

   Вывод:

   process 0 has [ 1.  2.  3.]
   process 1 has [ 4.  5.  6.]
   process 2 has [ 7.  8.  9.]
   

2. Разброс столбцов (не работает, все еще разбрасывает строки...):

   comm = MPI.COMM_WORLD
   rank = comm.Get_rank()
   size = comm.Get_size()
   
   A = np.zeros((3,3))
   if rank==0:
       A = np.matrix([[1.,2.,3.],[4.,5.,6.],[7.,8.,9.]]).T
   
   local_a = np.zeros(3)
   
   comm.Scatter(A, local_a, root=0)
   print "process", rank, "has", local_a
   

   Вывод:

   process 0 has [ 1.  2.  3.]
   process 1 has [ 4.  5.  6.]
   process 2 has [ 7.  8.  9.]
   

3. Разброс столбцов (работает, но кажется бессмысленным):

   comm = MPI.COMM_WORLD
   rank = comm.Get_rank()
   size = comm.Get_size()
   
   A = np.zeros((3,3))
   if rank==0:
       A = np.matrix([[1.,2.,3.],[4.,5.,6.],[7.,8.,9.]]).T.copy()
   
   local_a = np.zeros(3)
   
   comm.Scatter(A, local_a, root=0)
   print "process", rank, "has", local_a
   

   Наконец, давая желаемый результат:

   process 0 has [ 1.  4.  7.]
   process 2 has [ 3.  6.  9.]
   process 1 has [ 2.  5.  8.]
   

Есть ли простой способ отправить столбцы без копирования всей матрицы?

Для контекста я выполняю упражнение 5 в mpi4py учебнике. Мое полное решение (которое тратит память, как в пункте 3. выше) таково, если вы хотите знать:

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

A = np.zeros((3,3))
v = np.zeros(3)
result = np.zeros(3)
if rank==0:
    A = np.array([[1.,2.,3.],[4.,5.,6.],[7.,8.,9.]]).T.copy()
    v = np.array([0.1,0.01,0.001])

# Scatter the columns of the matrix
local_a = np.zeros(3)
comm.Scatter(A, local_a, root=0)

# Scatter the elements of the vector
local_v = np.array([0.])
comm.Scatter(v, local_v, root=0)

print "process", rank, "has A_ij =", local_a, "and v_i", local_v

# Multiplication
local_result = local_a * local_v

# Add together
comm.Reduce(local_result, result, op=MPI.SUM)
print "process", rank, "finds", result, "(", local_result, ")"

if (rank==0):
    print "The resulting vector is"
    print "   ", result, "computed in parallel"
    print "and", np.dot(A.T,v), "computed serially."

Вот тест профилирования памяти, запрошенный @Sajid:

Мое решение 3 (дает правильный ответ): 0.027 MiB A = np.array([[1.,2.,3.],[4.,5.,6.],[7.,8.,9.]]).T.copy() 0.066 MiB comm.Scatter(A, local_a, root=0) Всего = 0,093 МБ

Другое похожее решение (дает правильный ответ): 0.004 MiB A = np.array([[1.,2.,3.],[4.,5.,6.],[7.,8.,9.]]) 0.090 MiB comm.Scatter(A.T.copy(), local_a, root=0) Всего = 0,094 МБ.

@ Решение Саджида (дает правильный ответ): 0.039 MiB A[:,:] = np.transpose(np.array([[1.,2.,3.],[4.,5.,6.],[7.,8.,9.]])) 0.062 MiB comm.Scatter(A, local_a, root=0) Всего = 0,101 МБ

Мое решение 2 (дает неправильный ответ): 0.004 MiB A = np.array([[1.,2.,3.],[4.,5.,6.],[7.,8.,9.]]) 0.066 MiB comm.Scatter(A, local_a, root=0) Всего = 0,070 МБ

(Я только скопировал приращения памяти из строк, где приращение памяти различается между версиями кода. Очевидно, это все из корневого узла.)

Кажется очевидным, что все правильные решения должны копировать массив в память. Это неоптимально, так как все, что я хочу, это разбрасывать столбцы вместо строк.