Joe Todd
2020-12-01 21:10:27
我认为这段代码可以满足你的需求。这里的问题是Scatterv根本不关心numpy数组的形状，它只考虑包含值的线性内存块。因此，最简单的方法是事先将数据处理为正确的顺序。请注意，这send_data是一维数组，但这并不重要，因为Scatterv不在乎。在另一端，已经定义了recvbuf的形状，Scatterv只是从接收到的1D输入中填充它。
# Imports
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
rows = 4
num_cols = rows/size
send_data=None
if rank == 0:
data = np.matrix([[1, 2, 3, 4],
[5, 6, 7, 8],
[9, 10, 11, 12],
[13, 14, 15, 16]])
# Split into sub-arrays along required axis
arrs = np.split(data, size, axis=1)
# Flatten the sub-arrays
raveled = [np.ravel(arr) for arr in arrs]
# Join them back up into a 1D array
send_data = np.concatenate(raveled)
recvbuf = np.empty((rows, int(num_cols)), dtype='int')
comm.Scatterv(send_data, recvbuf, root=0)
print('Rank: ',rank, ', recvbuf received:\n ',recvbuf)
热门帖子
热门github
1
2
3
5
6
7
8
11
13
14
15
Mamba is a new state space model architecture showing promising performance on information-dense data such as language modeling, where previous subquadratic models fall short of Transformers. It is based on the line of progress on structured state space models, with an efficient hardware-aware design and implementation in the spirit of FlashAttention.
(翻译：Mamba 是一种新的状态空间模型架构，在信息密集型数据（例如语言建模）上显示出良好的性能，而之前的二次模型在 Transformers 方面存在不足。它基于结构化状态空间模型的进展，并本着FlashAttention的精神进行高效的硬件感知设计和实现。)
非常感谢您，您的回答很完美！