mindscience.e3nn.nn.Scatter

class mindscience.e3nn.nn.Scatter(mode='add')[source]

Easy-to-use wrapper for scatter operations: aggregates source values into a destination tensor according to index.

Parameters

mode (str, optional) –

{'add', 'sum', 'div', 'max', 'min', 'mul'}, scatter mode.

'add' or 'sum': element-wise addition.
'div': element-wise division.
'max': element-wise maximum.
'min': element-wise minimum.
'mul': element-wise multiplication.

Default: 'add'.

Inputs:

src (Tensor) - The source tensor to scatter.
index (Tensor) - The indices of elements to scatter, must be int type.
out (Tensor, optional) - The destination tensor. If provided, scatter will be performed in-place. Default: None.
dim_size (int, optional) - If out is not given, automatically create output with size dim_size. If dim_size is not given, a minimal sized output tensor is returned. Default: None.

Outputs:

output (Tensor) - The result after scatter operation.

Raises: ValueError – If mode is not legal.

Examples

>>> import mindspore as ms
>>> from mindspore import Tensor
>>> from mindscience.e3nn.nn import Scatter
>>> scatter = Scatter('add')
>>> src = Tensor([[1, 2], [3, 4], [5, 6]], ms.float32)
>>> index = Tensor([0, 0, 1], ms.int32)
>>> out = scatter(src, index)
>>> print(out.shape)
(3,2)