IndexAddOperation

功能

固定维度的指定下标加上某个特定值。

图1 IndexAddOperation

定义

struct IndexAddParam {
    enum IndexType {
        INDEX_UNDEFINED = 0,
        INDEX_ADD,
        INDEX_ADD_VALID, 
    };
    IndexType indexType = INDEX_UNDEFINED;
    int64_t axis = 0;
};

参数列表

成员名称	类型	默认值	描述
indexType	IndexType	INDEX_UNDEFINED	操作类型。 INDEX_UNDEFINED：未定义。 INDEX_ADD：加。 INDEX_ADD_VALID：有效长度内加。不支持Atlas 推理系列产品。
axis	int64_t	0	输入tensor需加上updates更新值的轴。当“indexType”为INDEX_ADD时，可为负数，取值范围为[-varDimNum, varDimNum - 1]。varDimNum为输入tensor0的维度数。当“indexType”为INDEX_ADD_VALID时，仅支持取值为0。

成员名称

类型

默认值

描述

indexType

IndexType

INDEX_UNDEFINED

操作类型。

INDEX_UNDEFINED：未定义。
INDEX_ADD：加。
INDEX_ADD_VALID：有效长度内加。不支持Atlas 推理系列产品。

axis

int64_t

输入tensor需加上updates更新值的轴。

当“indexType”为INDEX_ADD时，可为负数，取值范围为[-varDimNum, varDimNum - 1]。varDimNum为输入tensor0的维度数。
当“indexType”为INDEX_ADD_VALID时，仅支持取值为0。

输入输出描述（INDEX_ADD）

参数	维度	数据类型	格式	描述
var	[d_0, ..., d_n]	float16	ND	输入tensor。被加数，输入为零，原地被加后作为输出。
indices	[d_x]	int32	ND	输入tensor。指定固定维度的指定下标。d_min = min(d_x, d_axis)，值域为[0, d_min)，且前d_min个值不重复。
updates	[d_0, ..., d_n]	float16	ND	输入tensor。加数，根据indices的值加到var对应位置。维度数与var一致。索引为axis的维度为d_x，即d_axis == d_x。
alpha	[1]	float16	ND	输入tensor。累加次数。
output	[d_0, ..., d_n]	float16	ND	输出tensor。与var为同一个Tensor，即二者数据类型、数据格式和地址等所有属性均相同。
min(x, y)表示取x和y两个数中的较小值。 d_axis表示var、updates和output在索引为axis的维度大小，即：若axis为0，则d_axis对应第0维大小d_0。

输入输出描述（INDEX_ADD_VALID）

参数	维度	数据类型	格式	描述
var	[d_1, d_2]	float16	ND	输入tensor。被加数，输入为零，原地被加后作为输出。
indices	[d_0]	int32	ND	输入tensor。指定固定维度的指定下标。指定固定维度的指定下标。值域范围为[0, d_1)。
updates	[d_0, d_2]	float16	ND	输入tensor。加数，根据indices的值加到var对应位置。
validIndicesNum	[1]	int32	ND	输入tensor。indices的有效长度。值域范围为[0, d_0]。
output	[d_1, d_2]	float16	ND	输出tensor。与var为同一个Tensor，即二者数据类型、数据格式和地址等所有属性均相同。
d_2取值范围为(0, 8192]。

接口调用示例

indexType = INDEX_ADD

输入

indexType = 0

axis = 0

var =

[[ 1, 1, 1 ],
 [ 1，1，1],
 [ 1，1，1],
 [ 1，1，1],
 [ 1，1，1]]

indices = torch.tensor[0, 4, 2]

updates =

[[ 1, 2, 3 ],
 [ 4, 5, 6],
 [ 7, 8, 9]]

alpha = 1

输出

output =

[[  2,   3,   4],
 [  1,   1,   1],
 [  8,   9,  10],
 [  1,   1,   1],
 [  5,   6,   7]]

indexType = INDEX_ADD_VALID

输入

indexType = 1

axis = 0

var =

[[ 0, 0, 0 ],
 [ 0，0，0],
 [ 0，0，0],
 [ 0，0，0],
 [ 0，0，0]]

indices = {[0, 2, 2, -1, -1]}

updates =

[[ 0, 1, 2 ],
 [ 1, 2, 3],
 [ 2, 3, 4],
 [5, 6, 7],
 [6, 7, 8]]

validIndicesNum = {[3]}

输出

output =

[[  0,   1,   2],
 [  0,   0,   0],
 [  3,   5,  7],
 [  0,   0,   0],
 [  0,   0,   0]]

父主题： atb/infer_op_params.h