IndexAddOperation（部分代码开放）

产品支持情况

硬件型号	是否支持
Atlas A3 推理系列产品/Atlas A3 训练系列产品	√
Atlas A2 训练系列产品/Atlas 800I A2 推理产品	√
Atlas 训练系列产品	x
Atlas 推理系列产品	x
Atlas 200I/500 A2 推理产品	x

功能

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

图1 IndexAddOperation

定义

struct IndexAddParam {
    enum IndexType {
        INDEX_UNDEFINED = 0,
        INDEX_ADD,
        INDEX_ADD_VALID, 
    };
    IndexType indexType = INDEX_UNDEFINED;
    int64_t axis = 0;
    uint8_t rsv[16] = {0};
};

参数列表

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

成员名称

类型

默认值

取值范围

是否必选

描述

indexType

IndexType

INDEX_UNDEFINED

0，1，2

是

操作类型。

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

axis

int64_t

indexType = 1时：[-intensor[0].dimNum, intensor[0].dimNum - 1]
indexType = 2时：0

否

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

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

rsv[16]

uint8_t

{0}

[0]

否

预留参数。

输入输出描述（INDEX_ADD）

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

输入输出描述（INDEX_ADD_VALID）

参数	维度	数据类型	格式	描述
var	[dim_1, dim_2]	float16/bf16	ND	输入tensor。被加数，输入为零，原地被加后作为输出。
indices	[dim_0]	int32	ND	输入tensor。指定固定维度的指定下标。值域范围为[0, dim_1)。
updates	[dim_0, dim_2]	float16/bf16	ND	输入tensor。加数，根据indices的值加到var对应位置。
validIndicesNum	[1]	int32	ND	输入tensor。indices的有效长度。值域范围为[0, dim_0]。
output	[dim_1, dim_2]	float16/bf16	ND	输出tensor。与var为同一个Tensor，即二者数据类型、数据格式和地址等所有属性均相同。
dim_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