mindspore.ops.tensor_scatter_elements

View Source On Gitee
mindspore.ops.tensor_scatter_elements(input_x, indices, updates, axis=0, reduction='none')[source]

Write all elements in updates to the index specified by indices in input_x according to the reduction operation specified by reduction. axis controls the direction of the scatter operation.

tensor_scatter_elements takes three inputs input_x, updates and indices of the same rank r >= 1.

For a 3-D tensor, the output is:

output[indices[i][j][k]][j][k] = updates[i][j][k]  # if axis == 0, reduction == "none"

output[i][indices[i][j][k]][k] += updates[i][j][k]  # if axis == 1, reduction == "add"

output[i][j][indices[i][j][k]] = updates[i][j][k]  # if axis == 2, reduction == "none"

Warning

  • The order in which updates are applied is nondeterministic, meaning that if there are multiple index vectors in indices that correspond to the same position, the value of that position in the output will be nondeterministic.

  • On Ascend, the reduction only support set to "none" for now.

  • On Ascend, the data type of input_x must be float16 or float32.

  • This is an experimental API that is subject to change or deletion.

Note

If some values of the indices exceed the upper or lower bounds of the index of input_x, instead of raising an index error, the corresponding updates will not be updated to input_x. The backward is supported only for the case updates.shape == indices.shape.

Parameters

  • input_x (Tensor) – The target tensor. The rank must be at least 1.

  • indices (Tensor) – The index of input_x to do scatter operation whose data type must be mindspore.int32 or mindspore.int64. Same rank as input_x. And accepted range is [-s, s) where s is the size along axis.

  • updates (Tensor) – The tensor doing the scatter operation with input_x.

  • axis (int) – Which axis to scatter. Accepted range is [-r, r) where r = rank(input_x). Default: 0.

  • reduction (str) – Which reduction operation to scatter, supports "none" , "add" . Default: "none". When reduction is set to "none", updates will be assigned to input_x according to indices. When reduction is set to "add", updates will be added to input_x according to indices.

Returns

Tensor, has the same shape and type as input_x.

Raises

  • TypeError – If indices is neither int32 nor int64.

  • ValueError – If anyone of the rank among input_x, indices and updates less than 1.

  • ValueError – If the rank of updates is not equal to the rank of input_x.

  • RuntimeError – If the data type of input_x and updates conversion of Parameter is required when data type conversion of Parameter is not supported.

Supported Platforms:

Ascend GPU CPU

Examples

>>> import mindspore
>>> from mindspore import Tensor, ops
>>> from mindspore import Parameter
>>> import numpy as np
>>> input_x = Parameter(Tensor(np.array([[1, 2, 3, 4, 5]]), mindspore.int32), name="x")
>>> indices = Tensor(np.array([[2, 4]]), mindspore.int32)
>>> updates = Tensor(np.array([[8, 8]]), mindspore.int32)
>>> axis = 1
>>> reduction = "none"
>>> output = ops.tensor_scatter_elements(input_x, indices, updates, axis, reduction)
>>> print(output)
[[1 2 8 4 8]]
>>> input_x = Parameter(Tensor(np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]), mindspore.int32), name="x")
>>> indices = Tensor(np.array([[1, -1, 2], [0, 2, 1]]), mindspore.int32)
>>> updates = Tensor(np.array([[1, 2, 2], [4, 5, 8]]), mindspore.int32)
>>> axis = 0
>>> reduction = "add"
>>> output = ops.tensor_scatter_elements(input_x, indices, updates, axis, reduction)
>>> print(output)
[[ 5  2  3]
 [ 5  5 14]
 [ 7 15 11]]