megengine.functional.scatter¶

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

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

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

For a 3-D tensor, input tensor is updated as:

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
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]]

megengine.functional.gather

megengine.functional.cond_take