MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / scatter

Function scatter

imperative/python/megengine/functional/tensor.py:732–822  ·  view source on GitHub ↗

r""" Writes all values from the tensor source into input tensor at the indices specified in the index tensor. For each value in source, its output index is specified by its index in source for ``axis != dimension`` and by the corresponding value in index for ``axis = dimension``

(inp: Tensor, axis: int, index: Tensor, source: Tensor)

Source from the content-addressed store, hash-verified

730
731
732def scatter(inp: Tensor, axis: int, index: Tensor, source: Tensor) -> Tensor:
733 # TODO: rewrite doc
734 r"""
735 Writes all values from the tensor source into input tensor
736 at the indices specified in the index tensor.
737
738 For each value in source, its output index is specified by its index
739 in source for ``axis != dimension`` and by the corresponding value in
740 index for ``axis = dimension``.
741
742 For a 3-D tensor, input tensor is updated as:
743
744 .. code-block::
745
746 inp[index[i][j][k]][j][k] = source[i][j][k] # if axis == 0
747 inp[i][index[i][j][k]][k] = source[i][j][k] # if axis == 1
748 inp[i][j][index[i][j][k]] = source[i][j][k] # if axis == 2
749
750 ``inp``, ``index`` and ``source`` should have same number of dimensions.
751
752 It is also required that ``source.shape(d) <= inp.shape(d)`` and ``index.shape(d) == source.shape(d)``
753 for all dimensions ``d``.
754
755 Moreover, the values of index must be between ``0`` and ``inp.shape(axis) - 1`` inclusive.
756
757 Note:
758 Please notice that, due to performance issues, the result is uncertain on the GPU device
759 if scattering different positions from source to the same destination position
760 regard to index tensor.
761
762 Check the following examples, the oup[0][2] is maybe
763 from source[0][2] which value is 0.2256 or source[1][2] which value is 0.5339
764 if set the index[1][2] from 1 to 0.
765
766 Args:
767 inp: inp tensor which to be scattered.
768 axis: axis along which to index.
769 index: indices of elements to scatter.
770 source: source element(s) to scatter.
771
772 Return:
773 output tensor.
774
775 Examples:
776 >>> import numpy as np
777 >>> inp = Tensor(np.zeros(shape=(3,5),dtype=np.float32))
778 >>> source = Tensor([[0.9935,0.9465,0.2256,0.8926,0.4396],[0.7723,0.0718,0.5939,0.357,0.4576]])
779 >>> index = Tensor([[0,2,0,2,1],[2,0,1,1,2]])
780 >>> oup = F.scatter(inp, 0, index, source)
781 >>> oup.numpy()
782 array([[0.9935, 0.0718, 0.2256, 0. , 0. ],
783 [0. , 0. , 0.5939, 0.357 , 0.4396],
784 [0.7723, 0.9465, 0. , 0.8926, 0.4576]], dtype=float32)
785 """
786 input_shape = inp.shape
787 index_shape = index.shape
788 source_shape = source.shape
789 input_dims = len(input_shape)

Callers 2

funcFunction · 0.90
_select_and_scatterFunction · 0.50

Calls 3

_get_idxFunction · 0.85
formatMethod · 0.45
flattenMethod · 0.45

Tested by 1

funcFunction · 0.72