mindspore.ops.ReduceScatter
- class mindspore.ops.ReduceScatter(op=ReduceOp.SUM, group=GlobalComm.WORLD_COMM_GROUP)[source]
Reduces and scatters tensors from the specified communication group.
Note
The tensors must have the same shape and format in all processes of the collection.
- Parameters
- Inputs:
input_x (Tensor) - Input Tensor, suppose it has a shape \((N, *)\), where * means any number of additional dimensions. N must be divisible by rank_size. rank_size refers to the number of cards in the communication group.
- Outputs:
Tensor, it has the same dtype as input_x with a shape of \((N/rank\_size, *)\).
- Raises
TypeError – If any of operation and group is not a string.
ValueError – If the first dimension of the input cannot be divided by the rank_size.
- Supported Platforms:
Ascend
GPU
Examples
Note
Before running the following examples, you need to configure the communication environment variables.
For the Ascend devices, users need to prepare the rank table, set rank_id and device_id. Please see the rank table Startup for more details.
For the GPU devices, users need to prepare the host file and mpi, please see the mpirun Startup .
For the CPU device, users need to write a dynamic cluster startup script, please see the Dynamic Cluster Startup .
This example should be run with 2 devices.
>>> import mindspore as ms >>> from mindspore import Tensor >>> from mindspore.communication import init >>> from mindspore.ops import ReduceOp >>> import mindspore.nn as nn >>> import mindspore.ops as ops >>> import numpy as np >>> >>> ms.set_context(mode=ms.GRAPH_MODE) >>> init() >>> class Net(nn.Cell): ... def __init__(self): ... super(Net, self).__init__() ... self.reducescatter = ops.ReduceScatter(ReduceOp.SUM) ... ... def construct(self, x): ... return self.reducescatter(x) ... >>> input_ = Tensor(np.ones([8, 8]).astype(np.float32)) >>> net = Net() >>> output = net(input_) >>> print(output) [[2. 2. 2. 2. 2. 2. 2. 2.] [2. 2. 2. 2. 2. 2. 2. 2.] [2. 2. 2. 2. 2. 2. 2. 2.] [2. 2. 2. 2. 2. 2. 2. 2.]]
- Tutorial Examples: