megengine.functional.scatter

scatter(inp, axis, index, source)[源代码]

把张量 source 中所有的值通过 index 张量指定的索引位置上写入到输入张量中。

对于 source 中的每个值,它的输出索引在 axis != dimension 时,为 source 的索引或在 axis = dimension 时,为 index 中相对应的值。

对于三维张量, inp 将更新为:

inp[index[i][j][k]][j][k] = source[i][j][k]  # if axis == 0
inp[i][index[i][j][k]][k] = source[i][j][k]  # if axis == 1
inp[i][j][index[i][j][k]] = source[i][j][k]  # if axis == 2

inpindexsource 应当具有相同的维数。

在所有维度上需要满足 source.shape(d) <= inp.shape(d) 以及 index.shape(d) == source.shape(d)

此外, index 的值必须介于 0inp.shape(axis) - 1 之间(包含边界)。

注解

请注意,在GPU设备上,由于性能原因,若多个源数据被 index 指定同一个目标位置时,结果会不确定。

使用以下例子展示案例,如果将index[1][2]设置为1到0,则oup[0][2]可能来自值为0.2256的source[0][2],或值为0.5339的source[1][2]。

参数
  • inp – 将要进行scatter操作的inp张量。

  • axis – 将要进行索引的轴。

  • index – 将要进行scatter操作的元素的索引。

  • source – 将要进行scatter操作的inp张量一个或多个源元素。

返回

输出张量。

例如:

import numpy as np
import megengine.functional as F
from megengine import tensor

inp = tensor(np.zeros(shape=(3,5),dtype=np.float32))
source = tensor([[0.9935,0.9465,0.2256,0.8926,0.4396],[0.7723,0.0718,0.5939,0.357,0.4576]])
index = tensor([[0,2,0,2,1],[2,0,1,1,2]])
oup = F.scatter(inp, 0, index,source)
print(oup.numpy())

输出:

[[0.9935 0.0718 0.2256 0.     0.    ]
 [0.     0.     0.5939 0.357  0.4396]
 [0.7723 0.9465 0.     0.8926 0.4576]]