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
inp，index和source应当具有相同的维数。在所有维度上需要满足
source.shape(d) <= inp.shape(d)以及index.shape(d) == source.shape(d)。此外，
index的值必须介于0和inp.shape(axis) - 1之间(包含边界)。注解
请注意，在GPU设备上，由于性能原因，若多个源数据被 index 指定同一个目标位置时，结果会不确定。
使用以下例子展示案例，如果将index[1][2]设置为1到0，则oup[0][2]可能来自值为0.2256的source[0][2]，或值为0.5339的source[1][2]。
- 参数
inp (
Tensor) – 将要进行scatter操作的inp张量。axis (
int) – 将要进行索引的轴。index (
Tensor) – 将要进行scatter操作的元素的索引。source (
Tensor) – 将要进行scatter操作的inp张量一个或多个源元素。
- 返回类型
Tensor- 返回
输出张量。
例如：
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]]