megengine.functional.scatter¶

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

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

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

3 维的输入 Tensor 将会按照如下方式更新：

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 操作的一个或多个元素。

返回类型

Tensor

返回

输出张量。

实际案例

>>> import numpy as np
>>> 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)
>>> oup.numpy()
array([[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]], dtype=float32)

megengine.functional.gather

megengine.functional.cond_take